I had a similar issue and I think it is related to your optimizer referencing model parameters that are not yet loaded on GPU. In my case I fixed the issue by changing the following:
optimizer = torch.optim.Adam(model.parameters(), lr=params.lr)
try:
       checkpoint = torch.load("checkpoints/"+model_folder+"checkpoint.pth.tar")
       start_epoch = checkpoint['epoch']
       scheduler.load_state_dict(checkpoint['scheduler'])
       model.load_state_dict(checkpoint['state_dict'])
       model.to(device) # moving this line here fixed the issue
       optimizer.load_state_dict(checkpoint['optimizer'])
       …
       # it used to be somewhere around here

So in your case a cleaner solution would be to send that parameters object to GPU (not sure how but there must be some one-liner that does that) BEFORE loading the dictionary, else it seems that it will be loaded to the usual memory instead of GPU memory.