Hi. I checked the parameters device and all of them seem to be on cpu.
I tried to check if the gpu is active using - print(torch.cuda.is_available()) and it shows True.

Is there any other way by which we can assign the device to be GPU (for each parameter)?
TIA
Edit: Hi. I found the error. I had to include the following lines of code before loading the trained parameters into the new model.
    device = torch.device('cuda')
    custom_model = Net().to(device)
    custom_model.state_dict(torch.load(PATH))

Now all parameters are on the GPU. Thanks a lot!