You could “remove” nn.Modules by replacing them with nn.Identity.
However, this won’t work for functional calls in your model, since (as you already explained), these methods are not registered as modules, so you would need to override the forward.