Solved by gfrogat in post #2 
              
              
                Hi 
Based on your error message, I guess that the data-types of the tensors in the momentum_buffer of the optimizer are on the wrong device. 

You can try to run this snippet to move the buffers: 
print(optimizer.state[list(optimizer.state.keys())[0]])

for p in optimizer.state.keys():
    param_sta…