Hi,
You shouldn’t use the Module.to() functions inside your forward. It won’t work well with AD.
Only the Tensor.to() that you were mentioning above will.
The reason for this is that the Module Parameters are always leaf Tensor. And so cannot be updated in a differentiable manner.
Note that more generally, I don’t think your approach will help wrt to GPU memory as the autograd needs to keep values to be able to compute the backward pass. And your model’s weights are part of them. So even if you move the model to the cpu, these buffers will stay on the GPU.