nn.DataParallel: RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!

Massimiliano_Datres: <p>Dear all,</p>
<p>I am trying to parallelise my model to train it on 2 GPUs. I have decided to use the nn.DataParallel function, but I received the following error during the forward pass:</p>
<blockquote>
<p>Traceback (most recent call last):<br>
File “main.py”, line 93, in <br>
args.func(args)<br>
File “/home/mdatres/new_quantlab/quantlab/manager/flows/train.py”, line 243, in train<br>
ypr = net(x)<br>
File “/data2/mdatres/miniconda3/envs/quantlab/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1051, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/data2/mdatres/miniconda3/envs/quantlab/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py”, line 168, in forward<br>
outputs = self.parallel_apply(replicas, inputs, kwargs)<br>
File “/data2/mdatres/miniconda3/envs/quantlab/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py”, line 178, in parallel_apply<br>
return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])<br>
File “/data2/mdatres/miniconda3/envs/quantlab/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py”, line 86, in parallel_apply<br>
output.reraise()<br>
File “/data2/mdatres/miniconda3/envs/quantlab/lib/python3.8/site-packages/torch/_utils.py”, line 425, in reraise<br>
raise self.exc_type(msg)<br>
RuntimeError: Caught RuntimeError in replica 1 on device 1.<br>
Original Traceback (most recent call last):<br>
File “/data2/mdatres/miniconda3/envs/quantlab/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py”, line 61, in _worker<br>
output = module(*input, **kwargs)<br>
File “/data2/mdatres/miniconda3/envs/quantlab/lib/python3.8/site-packages/torch/fx/graph_module.py”, line 513, in wrapped_call<br>
raise e.with_traceback(None)<br>
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!</p>
</blockquote>
<p>I used the resnet50 model which is already implemented in torchvision.</p>
<pre><code class="lang-auto">class ResNet(nn.Module):
    
    def __init__(self, num_classes = 2, pretrained = True, seed = -1):
        super().__init__()
        
        self.net1 = resnet50(pretrained=pretrained)
        self.net1.fc = nn.Linear(in_features=2048, out_features=num_classes, bias=True)
        self.softmax = nn.Softmax()
        # if not pretrained:
        #      self.net1._initialize_weights(seed=seed)



    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        
        x = self.net1(x)

        return self.softmax(x)

</code></pre>
<p>Below I attach the code that I have used to load the model on the GPUs and to train it:</p>
<pre><code class="lang-auto">  # cycle over epochs (one loop for each fold)
        for epoch_id in range(start_epoch_id, logbook.n_epochs):

            # === EPOCH: START ===

            # === TRAINING STEP: START ===
            
            net.train()
            
            # master-workers synchronisation point: quantization controllers might change the network's quantization parameters stochastically
            if (not platform.is_horovod_run) or platform.is_master:
                for c in qnt_ctrls:
                    c.step_pre_training_epoch(epoch_id)
                
            if platform.is_horovod_run:
                platform.hvd.broadcast_parameters(net.state_dict(), root_rank=platform.master_rank)

            # cycle over batches of training data (one loop for each epoch)
            for batch_id, (x, ygt) in enumerate(train_loader):

                # master-workers synchronisation point: quantization controllers might change the network's quantization parameters stochastically
                # TODO: in multi-process runs, synchronising processes at each step might be too costly
                if (not platform.is_horovod_run) or platform.is_master:
                    for c in qnt_ctrls:
                        c.step_pre_training_batch()
                if platform.is_horovod_run:
                    platform.hvd.broadcast_parameters(net.state_dict(), root_rank=platform.master_rank)

                # event: forward pass is beginning
                train_meter.step(epoch_id, batch_id)
                train_meter.start_observing()
                train_meter.tic()

                # processing (forward pass)
                x   = x.to(platform.device)
                
                ypr = net(x)

                # loss evaluation
                ygt  = ygt.to(platform.device)
                loss = loss_fn(ypr, ygt)

                # event: forward pass has ended; backward pass is beginning
                train_meter.update(ygt, ypr, loss)

                # training (backward pass)
                gd.opt.zero_grad()  # clear gradients
                loss.backward()     # gradient computation
                gd.opt.step()       # gradient descent

                # event: backward pass has ended
                train_meter.toc(ygt)
                train_meter.stop_observing()
</code></pre>
<p>In what’s follow the variable platform.device is “cuda:0”.</p>
<pre><code class="lang-auto">net = net.to(platform.device)

        if platform.is_nndataparallel_run:
            net = nn.DataParallel(net)  # single-node, single-process, multi-GPU run
            
</code></pre>
<p>Thanks for your help.</p>

ptrblck: <p>I cannot reproduce the issue and your model works fine:</p>
<pre><code class="lang-python">import torch.nn as nn
import torchvision.models as models

class ResNet(nn.Module):
    def __init__(self, num_classes = 2, pretrained = True, seed = -1):
        super().__init__()
        self.net1 = models.resnet50(pretrained=pretrained)
        self.net1.fc = nn.Linear(in_features=2048, out_features=num_classes, bias=True)
        self.softmax = nn.Softmax()

    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        print('inside forward, x.device {}'.format(x.device))
        x = self.net1(x)
        return self.softmax(x)

model = ResNet()
model = nn.DataParallel(model).cuda()
x = torch.randn(8, 3, 224, 224).cuda()

out = model(x)
print(out.device)
</code></pre>
<p>Output:</p>
<pre><code class="lang-python">inside forward, x.device cuda:0
inside forward, x.device cuda:1
inside forward, x.device cuda:2
inside forward, x.device cuda:3
inside forward, x.device cuda:4
inside forward, x.device cuda:5
inside forward, x.device cuda:6
inside forward, x.device cuda:7
tmp.py:24: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.
  return self.softmax(x)
cuda:0
</code></pre>

Massimiliano_Datres: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="158916">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p><code>x = torch.randn(8, 3, 224, 224).cuda()</code></p>
</blockquote>
</aside>
<p>Thanks for the quick reply. The problem was not in the model but in the step_pre_training_epoch. Returning the model here solved my issue.</p>