DataParallel expecting all tensors in the same device

username: adria
cooked: <p>Hello, I am using DataParallel in a similar way as shown in this <a href="https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html" rel="noopener nofollow ugc">tutorial</a>.<br>
I run into the following error:</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 mat2 in method wrapper_mm)</p>
</blockquote>
<p>This happens right after the first forward pass.<br>
The model architecture is built from <a href="https://pytorch-geometric-temporal.readthedocs.io/en/latest/index.html" rel="noopener nofollow ugc">Pytorch Geometric temporal</a></p>
<pre data-code-wrap="python"><code class="lang-python">
#How the data is loaded
train_dataset, test_dataset = temporal_signal_split(dataset, train_ratio=0.7)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

#Model definition
class RecurrentGCN(torch.nn.Module):
    def __init__(self, node_features):
        super(RecurrentGCN, self).__init__()
        torch.manual_seed(1234567)
        #from Pytorch Geometric Temporal
        self.recurrent1 = SomeGCN(node_features, 32, 1)
        self.fc1 = torch.nn.Linear(32, 1)
        self.sigmoid = torch.sigmoid
        self.dropout = torch.nn.Dropout(0.2)
    
    def forward(self, data):
        x, edge_index, edge_weight = data.x, data.edge_index, data.edge_attr
        x = torch.flatten(x, start_dim=1)
        h_0 = self.recurrent1(x, edge_index, edge_weight)
        h = F.relu(h_0)
        h = self.fc1(h)
        h = self.sigmoid(h)
        h = self.dropout(h)
        return h 

#DataParallel instance
model = RecurrentGCN(node_features = n_features)
    if torch.cuda.device_count() &gt; 1:
        print("Available/CUDA_VISIBLE_DEVICES", os.environ["CUDA_VISIBLE_DEVICES"])
        print("Device count", torch.cuda.device_count())
        # dim = 0 [30, xxx] -&gt; [10, ...], [10, ...], [10, ...] on 3 GPUs
        model = DataParallel(model, device_ids=[0, 1])
    model.to(device)

#during train loop
for snapshot in train_dataset:
            snapshot = snapshot.to(device)
            y_hat = model(snapshot) #Error happens here
</code></pre>
<p>The tutorial seems straightforward so I do not see what am I doing wrong here.<br>
Last note, the code runs without problems in a single GPU.</p>
<p>Thanks for the help!</p>

username: ptrblck
cooked: <p>Could you check the <code>.device</code> attribute of all tensors inside the <code>forward</code> method and compare them to the <code>.device</code> of the parameters?<br>
They should match and based on your description I would guess that the passed <code>input</code> might be a custom class, which is not recognized by <code>nn.DataParallel</code> and thus not split into the chunks for each GPU.</p>

username: adria
cooked: <p>I think you are right. I put a breakpoint on the forward method and I can see for example fc1 and recurrent1 parameters are on cuda:1 whilst the attribute tensors are on cuda:0.<br>
How could I try to solve this?</p>

username: ptrblck
cooked: <p>Make sure that tensors are passed to the model, which can be split in their batch dimension.</p>

username: adria
cooked: <p>I have updated the training loop to send tensors to the model that can be split in their batch dimension like this:</p>
<pre data-code-wrap="python"><code class="lang-python">for snapshot in train_dataset:
            x, edge_index, edge_weight = snapshot.x, snapshot.edge_index, snapshot.edge_attr
            x = torch.flatten(x, start_dim=1).to(device)
            edge_index = edge_index.to(device)
            edge_weight = edge_weight.to(device)
            y_hat = model(x, edge_index, edge_weight)
</code></pre>
<p>This solves the previous error, but it results in a new one:</p>
<blockquote>
<p>Exception has occurred: IndexError       (note: full exception trace is shown but execution is paused at: _run_module_as_main)<br>
Caught IndexError in replica 0 on device 0.<br>
Original Traceback (most recent call last):<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/site-packages/torch/nn/parallel/parallel_apply.py”, line 61, in _worker<br>
output = module(*input, **kwargs)<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/site-packages/torch/nn/modules/module.py”, line 1102, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/mnt/raid0/users/acg384/workspace/code/LSTM_ddp2.py”, line 49, in forward<br>
h_0 = self.recurrent1(x, edge_index, edge_weight)<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/site-packages/torch/nn/modules/module.py”, line 1102, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/site-packages/torch_geometric_temporal/nn/recurrent/gconv_gru.py”, line 163, in forward<br>
Z = self._calculate_update_gate(X, edge_index, edge_weight, H)<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/site-packages/torch_geometric_temporal/nn/recurrent/gconv_gru.py”, line 120, in _calculate_update_gate<br>
Z = self.conv_x_z(X, edge_index, edge_weight)<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/site-packages/torch/nn/modules/module.py”, line 1102, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/site-packages/torch_geometric/nn/conv/cheb_conv.py”, line 143, in forward<br>
edge_index, norm = self.<strong>norm</strong>(edge_index, x.size(self.node_dim),<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/site-packages/torch_geometric/nn/conv/cheb_conv.py”, line 110, in <strong>norm</strong><br>
edge_index, edge_weight = remove_self_loops(edge_index, edge_weight)<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/site-packages/torch_geometric/utils/loop.py”, line 36, in remove_self_loops<br>
mask = edge_index[0] != edge_index[1]<br>
IndexError: index 1 is out of bounds for dimension 0 with size 1<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/site-packages/torch/_utils.py”, line 434, in reraise<br>
raise exception<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/site-packages/torch/nn/parallel/parallel_apply.py”, line 86, in parallel_apply<br>
output.reraise()<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/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 “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/site-packages/torch/nn/parallel/data_parallel.py”, line 168, in forward<br>
outputs = self.parallel_apply(replicas, inputs, kwargs)<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/site-packages/torch/nn/modules/module.py”, line 1102, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/mnt/raid0/users/acg384/workspace/code/LSTM_ddp2.py”, line 176, in <br>
y_hat = model(x, edge_index, edge_weight)<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/runpy.py”, line 87, in _run_code<br>
exec(code, run_globals)<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/runpy.py”, line 97, in _run_module_code<br>
_run_code(code, mod_globals, init_globals,<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/runpy.py”, line 268, in run_path<br>
return _run_module_code(code, init_globals, run_name,<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/runpy.py”, line 87, in _run_code<br>
exec(code, run_globals)<br>
File “/home/acg384/miniconda3/envs/pytorch_test/lib/python3.9/runpy.py”, line 197, in _run_module_as_main (Current frame)<br>
return _run_code(code, main_globals, None,</p>
</blockquote>
<p>This seems to be triggered by some internals within pytorch geometric.</p>