Strange issue.  I was using SGD on my model, and saved and loaded the optimiser state fine on GPU. But replacing SGD with Adam suddenly complains that it’s internal state requires using CPU (using Pytorch 1.2, CUDA 10):

Traceback (most recent call last):
File “/root/network/net.py”, line 139, in train
optimiser.step()
File “/opt/conda/lib/python3.6/site-packages/torch/optim/adam.py”, line 93, in step
exp_avg.mul_(beta1).add_(1 - beta1, grad)
RuntimeError: expected device cpu and dtype Float but got device cuda:0 and dtype Float

The very same code saves and loads SGD optimiser state without problems. Should I move and load Adam somewhat differently? This is the code that loads the state and trains:
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  model = MyModel()

  # optimiser init
  optimiser = torch.optim.Adam(model.parameters(), lr=0.1)
  first_epoch = 0

  if load_model != None:
    checkpoint = torch.load(load_model, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimiser.load_state_dict(checkpoint['optimiser_state_dict'])
    first_epoch = checkpoint['epoch'] + 1

  model.to(device)
  for epoch in range():
     optimiser.zero_grad()
     ...
     # training here
     model(x, y)
     ...
     loss.backward()
     optimiser.step() # this is where the error happens

This is how I save the state:
        torch.save({
              'epoch': epoch,
              'model_state_dict': model.state_dict(),
              'optimiser_state_dict': optimiser.state_dict(),
              'loss': epoch_loss,
              }, model_file)
        print("Model saved:", model_file)

I see from the error message that is is caused by internal parameters of Adam, but as I am already loading checkpoint to GPU, I cannot understand, why?