Suppose I have a model to train, and I also do some additional post processing on the output of the model. The post processing is parameterized by a scaler tensor which I also want to train. How should I tell the optimizer to train both the model and the scaler?
If your post processing method uses nn.Parameters
, you could simply pass them additionally to the optimizer as described here.
1 Like