Code that loads SGD fails to load Adam state to GPU

martinr: <p>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):</p>
<blockquote>
<p>Traceback (most recent call last):<br>
File “/root/network/net.py”, line 139, in train<br>
optimiser.step()<br>
File “/opt/conda/lib/python3.6/site-packages/torch/optim/adam.py”, line 93, in step<br>
exp_avg.mul_(beta1).add_(1 - beta1, grad)<br>
RuntimeError: expected device cpu and dtype Float but got device cuda:0 and dtype Float</p>
</blockquote>
<p>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:</p>
<pre><code class="lang-auto">  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
</code></pre>
<p>This is how I save the state:</p>
<pre><code class="lang-auto">        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)
</code></pre>
<p>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?</p>

ptrblck: <p>The internal states might have been stored as CUDATensors, if you’ve pushed the model to the GPU in your previous run.<br>
Does your code work, if you push the model to the GPU again before initializing the optimizer?</p>

Upgrade_Yourself: <p>I think you’re right.<br>
In <a href="http://pytorch.org" rel="noopener nofollow ugc">pytorch.org</a>, this note is mentioned:</p>
<blockquote>
<p>If you need to move a model to GPU via  <code>.cuda()</code> , please do so before constructing optimizers for it. Parameters of a model after  <code>.cuda()</code>  will be different objects with those before the call.<br>
In general, you should make sure that optimized parameters live in consistent locations when optimizers are constructed and used.</p>
</blockquote>

martinr: <p>Thank you! Indeed, pushing the model to GPU before loading the state of optimiser solved the problem. Though, I would say that the error message should pop up upon loading the state dict, not on step(). Thanks again <a class="mention" href="/u/ptrblck">@ptrblck</a> and <a class="mention" href="/u/upgrade_yourself">@Upgrade_Yourself</a>!</p>

Flock1: <p>I have created a model class like this:</p>
<pre><code class="lang-auto">class model_parallel(nn.Module):
   def __init__(self):
      super().__init__()
      sub_net_1 = models.resnet50(True)
      # sub_net_1 = torch.nn.Sequential(*(list(sub_net_1.children())[:-3]))
      for param in sub_net_1.parameters():
         param.requires_grad = False
      sub_net_2 = nn.Sequential(nn.Linear(in_features=1000, out_features=500,bias=True),
      nn.ReLU(),
      nn.Linear(in_features=500, out_features=100,bias=True),
      nn.ReLU(),
      nn.Linear(in_features=100, out_features=67,bias=True))
      self.sub_network1 = sub_net_1.cuda(0)
      self.sub_network2 = sub_net_2.cuda(1)

   def forward(self, x):
      x = x.cuda(0)
      x = self.sub_network1(x)
      # print(x.shape)
      x = x.cuda(1)
      x = self.sub_network2(x)
      return x
</code></pre>
<p>However, I am getting</p>
<blockquote>
<p>RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument mat1 in method wrapper_addmm)</p>
</blockquote>
<p>at</p>
<blockquote>
<p>x = self.sub_network2(x)</p>
</blockquote>
<p>What am I doing wrong here?</p>

ptrblck: <p>Your code works fine using:</p>
<pre><code class="lang-python">model = model_parallel()
x = torch.randn(64, 3, 224, 224, device='cuda')
out = model(x)
print(out.shape)
&gt; torch.Size([64, 67])
</code></pre>