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.



 kirtac83:

Traceback (most recent call last): File “/opt/anaconda3/lib/python3.7/site-packages/torch/optim/sgd.py”, line 101, in step buf.mul_(momentum).add_(1 - dampening, d_p) RuntimeError: expected type torch.FloatTensor but got torch.cuda.FloatTensor


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_state = optimizer.state[p]
    buf = param_state["momentum_buffer"]
    param_state["momentum_buffer"] = buf.cuda()  # move buf to device

print(optimizer.state[list(optimizer.state.keys())[0]])

See here 15 how momentum_buffers are used in the SGD optimizer.
Hope this helps.