I am facing a weird issue. I am doing the following


Freeze a torchvision model on CPU host, which internalizes parameters/weight tensors as constants.
Save this frozen module.
Try to load and run the model on GPU host.

I see that the weight tensor constants are not getting loaded onto GPU despite me invoking model.to(cuda:0) on the loaded model.
How do we ensure that the tensor constants, of the saved model, get loaded onto GPU?
Error StackTrace:
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/torchvision/models/segmentation/fcn/___torch_mangle_17301.py", line 9, in forward
    input_shape = torch.slice(torch.size(x), -2, 9223372036854775807, 1)
    features = torch.dict()
    x0 = torch.conv2d(x, CONSTANTS.c0, None, [2, 2], [3, 3], [1, 1], 1)
         ~~~~~~~~~~~~ <--- HERE
    x1 = torch.batch_norm(x0, CONSTANTS.c1, CONSTANTS.c2, CONSTANTS.c3, CONSTANTS.c4, False, 0.10000000000000001, 1.0000000000000001e-05, True)
    x2 = torch.relu_(x1)
 
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same