Pytorch / device problem(cpu, gpu) when load state dict for optimizer

Question

Hi i`m student who studies pytorch since last summer.

state = torch.load('drive/My Drive/MODEL/4 CBAM classifier55')

model = MyResNet()
model.load_state_dict(state['state_dict'])

criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters(), lr=0.0003,betas=(0.5,0.999))
optimizer.load_state_dict(state['optimizer'])

model.to(device)

i wrote code like above.

RuntimeError                              Traceback (most recent call last)
<ipython-input-26-507493db387a> in <module>()
     56     new_loss.backward()
     57 
---> 58     optimizer.step()
     59 
     60     running_loss += loss.item()

/usr/local/lib/python3.6/dist-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
     13         def decorate_context(*args, **kwargs):
     14             with self:
---> 15                 return func(*args, **kwargs)
     16         return decorate_context
     17 

/usr/local/lib/python3.6/dist-packages/torch/optim/adam.py in step(self, closure)
     97 
     98                 # Decay the first and second moment running average coefficient
---> 99                 exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
    100                 exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
    101                 if amsgrad:

RuntimeError: expected device cpu but got device cuda:0

And when i implement training code, then i got this kind of error. When i comment out 'optimizer.load_state_dict', it works well. How can i solve this problem? Thank you for your answer. :)

Answer 1

Seems like the state was on cuda when you saved and now trying to use it on cpu or vice-versa. To avoid this error, a simple way is to pass the map_location argument to load.

Just pass map_location=<device you want to use> in torch.load and it should work fine. Also, see https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-loading-model-across-devices