I’m not familiar with your use case and don’t know when you would like to use it.
If you want to change the forward
method after the training, you could use e.g. a flag you could pass to it and select the desired code path (with the manipulation) or use forward hooks as described here.