[resolved] Some of weight/gradient/input tensors are located on different GPUs Error

username: SelvamArul
cooked: <p>I am training a CNN on top of the features from a pretrained network for doing semantic segmentation. I have these models on two different GPUs and I copy the data to the respective GPUs as shown in the following snippet.</p>
<pre><code># move pretrained model from CPU to GPU 1
backbone = backbone.cuda(1)

# move the current model to GPU2
model = model.cuda(2)

# NLL criterion
criterion = nn.NLLLoss2d().cuda(2)

   # IN training loop
    x = Variable(t_rgb[i * 4 : i * 4 + cbs].type(dtype).cuda(1), requires_grad=False)
    y = Variable(t_target[i * 4 : i * 4 + cbs].type(th.LongTensor).cuda(2), requires_grad=False)
    
    x = backbone.forward(x).type(dtype).cuda(2)
    output = model(x)
    
    print ('x ', x.get_device())
    print ('output ',output.get_device())
    print ('y ',y.get_device())
    loss = criterion(output, y)
    print ('loss ',loss.get_device())
    optimizer.zero_grad()
    loss.backward()
</code></pre>
<p>The output I get is,<br>
x 2<br>
output  2<br>
y  2<br>
loss  2</p>
<p>But then I get</p>
<p><code>RuntimeError: Assertion</code>THCTensor_(checkGPU)(state, 3, input, gradInput, gradOutput)’ failed. Some of weight/gradient/input tensors are located on different GPUs. Please move them to a single one. at /b/wheel/pytorch-src/torch/lib/THCUNN/generic/Threshold.cu:49`</p>
<p>The exact stacktrace is as follows.</p>
<pre><code>  File "training.py", line 319, in &lt;module&gt;
    loss.backward()
  File "/opt/python3.5/lib/python3.5/site-packages/torch/autograd/variable.py", line 146, in backward
    self._execution_engine.run_backward((self,), (gradient,), retain_variables)
  File "/opt/python3.5/lib/python3.5/site-packages/torch/nn/_functions/thnn/auto.py", line 175, in backward
    update_grad_input_fn(self._backend.library_state, input, grad_output, grad_input, *gi_args)
RuntimeError: Assertion `THCTensor_(checkGPU)(state, 3, input, gradInput, gradOutput)' failed. Some of weight/gradient/input tensors are located on different GPUs. Please move them to a single one. at /b/wheel/pytorch-src/torch/lib/THCUNN/generic/Threshold.cu:49
</code></pre>
<p>From what I understand, input to the model (Variable x) and the model (and so the weights and the gradient) are both on the same GPU (2 in my case). But I don’t understand why I am getting this error. Could someone please help me in figuring out the issue here?  Also, I welcome any suggestions related to having models in different GPUs and moving data around.</p>

username: SelvamArul
cooked: <p>I managed to figure out the issue. I didn’t set <code>requires_grad</code> flag to <code>False</code> for the parameters in the pretrained network.</p>

username: Baichuan
cooked: <p>Hi, I have the same problem as you. I want to train the same model in two different GPU.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.optim as optim

class ToyModel(nn.Module):
  def __init__(self):
    super(ToyModel, self).__init__()
    self.net1 = torch.nn.Linear(10, 10)  
    self.relu = torch.nn.ReLU()
    self.net2 = torch.nn.Linear(10, 5)  

  def forward(self, x):
    x = self.relu(self.net1(x))
    return self.net2(x)
    
model = ToyModel()
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

optimizer.zero_grad()

outputs1 = model.cuda(1)(torch.randn(20, 10).cuda(1)) #gpu=1
outputs2 = model.cuda(2)(torch.randn(20,10).cuda(2)) #gpu=2
labels = torch.randn(20, 5).cuda(3) #gpu=3 loss
loss_fn(outputs1.cuda(3)+outputs2.cuda(3),labels).backward()
optimizer.step()
</code></pre>
<p>But the error is RuntimeError: binary_op(): expected both inputs to be on same device, but input a is on cuda:1 and in│<br>
put b is on cuda:2</p>

username: ptrblck
cooked: <p>The code looks alright.<br>
To debug further, I would recommend to split the operations and check the <code>.device</code> attribute of all tensors separately.</p>

username: Baichuan
cooked: <p>Thanks for your reply. In the above code,  outputs1, outputs2 and labels are on GPU3. The model are on GPU2. But I don’t know why it can’t work. The error imformation is “but input a is on cuda:1 and input b is on cuda:2”</p>

username: aeternum
cooked: <p>Lets get the relevant facts straight before we get into the issue:</p>
<ol>
<li>We are training the model on multiple GPUs.</li>
<li>We are moving the tensors to GPU using .cuda()</li>
</ol>
<p>The issue:<br>
While performing a computation it is required that the operands (here tensors) be on the same device.</p>
<ul>
<li>While computing the loss i.e criterion(labels, predicted) the labels and predicted values are on different GPUs.</li>
</ul>
<p>Why has this happened:<br>
The code has made use of .cuda() to migrate tensors from CPU to GPU which does not specify a particular device (eg.0, 1, 2 or 3), hence labels and predicted values may very well be on different GPUs.</p>
<p>Solution:<br>
Just as the the training loop begins, use <code>torch.cuda.set_device(&lt;device&gt;)</code>.</p>
<ul>
<li>A witty way of still utilising all GPUs would be : <code>torch.cuda.set_device(iteration % the_number_of_gpus)</code>.</li>
</ul>
<p>This will resolve the error.</p>