Train model parameters and scaler parameters together

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