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 it seems you need to send de model to GPU BEFORE loading the optimizer dictionary, else it will be loaded to the usual memory instead of GPU memory, which caused your issue (I think).