By “internal states” I mean all class attributes in the modules.
In particular buffers and parameters are of interest, as they would need to be pushed to the appropriate device (e.g. such as the weight buffer in nn.CrossEntropyLoss).
The dim argument in nn.Softmax however will not be pushed to the device, as it’s a plain Python integer to specify the dimension the softmax is applied in.