DataParallel Not working with timm Model

Sheikh_Talha: <p>Hello,</p>
<p>I am trying to train a model which is based on timm library.</p>
<p>My main looks as:<br>
model = DPTDepthModel(<br>
path=None,<span class="hashtag">#model_path</span>,<br>
scale=0.00006016,<br>
shift=0.00579,<br>
invert=True,<br>
backbone=“vitb_rn50_384”,<br>
non_negative=True,<br>
enable_attention_hooks=False,<br>
)</p>
<pre><code>loss_function = MSE()

optimizer = torch.optim.Adam(model.parameters(), lr=0.001, )

if torch.cuda.device_count() &gt; 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    torch.cuda.empty_cache()
    model = nn.DataParallel(model, device_ids = [0, 1])
    model.cuda()
    print("Model and Loss Function are on multiple GPUs now")

for epoch in range(0, 5):
    newmodel=train(model,train_loader,loss_function,optimizer,epoch)
    torch.save(newmodel.state_dict(), "model"+ str(epoch) +".pt")    
return
</code></pre>
<p>Error:</p>
<p>File “train.py”, line 482, in main2<br>
newmodel=train(model,train_loader,loss_function,optimizer,epoch)<br>
File “train.py”, line 381, in train<br>
output=model(input.float().cuda())<span class="hashtag">#to</span>(torch.device(‘cuda:1’)))<br>
File “/home/talha/anaconda3/envs/SemNet2/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 1102, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/talha/anaconda3/envs/SemNet2/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py”, line 168, in forward<br>
outputs = self.parallel_apply(replicas, inputs, kwargs)<br>
File “/home/talha/anaconda3/envs/SemNet2/lib/python3.6/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/talha/anaconda3/envs/SemNet2/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py”, line 86, in parallel_apply<br>
output.reraise()<br>
File “/home/talha/anaconda3/envs/SemNet2/lib/python3.6/site-packages/torch/_utils.py”, line 434, in reraise<br>
raise exception<br>
RuntimeError: Caught RuntimeError in replica 1 on device 1.<br>
Original Traceback (most recent call last):<br>
File “/home/talha/anaconda3/envs/SemNet2/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py”, line 61, in _worker<br>
output = module(*input, **kwargs)<br>
File “/home/talha/anaconda3/envs/SemNet2/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 1102, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/talha/SemAttNet-Thesis-Code-main/dpt/models.py”, line 165, in forward<br>
inv_depth = super().forward(x).squeeze(dim=1)<br>
File “/home/talha/SemAttNet-Thesis-Code-main/dpt/models.py”, line 78, in forward<br>
glob = self.pretrained.model.forward_flex(x)<br>
File “/home/talha/SemAttNet-Thesis-Code-main/dpt/vit.py”, line 175, in forward_flex<br>
x = self.patch_embed.backbone(x)<br>
File “/home/talha/anaconda3/envs/SemNet2/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 1102, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/talha/anaconda3/envs/SemNet2/lib/python3.6/site-packages/timm/models/resnetv2.py”, line 418, in forward<br>
x = self.forward_features(x)<br>
File “/home/talha/anaconda3/envs/SemNet2/lib/python3.6/site-packages/timm/models/resnetv2.py”, line 412, in forward_features<br>
x = self.stem(x)<br>
File “/home/talha/anaconda3/envs/SemNet2/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 1102, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/talha/anaconda3/envs/SemNet2/lib/python3.6/site-packages/torch/nn/modules/container.py”, line 141, in forward<br>
input = module(input)<br>
File “/home/talha/anaconda3/envs/SemNet2/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 1102, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/talha/anaconda3/envs/SemNet2/lib/python3.6/site-packages/timm/models/layers/std_conv.py”, line 72, in forward<br>
x = F.conv2d(x, weight, self.bias, self.stride, self.padding, self.dilation, self.groups)<br>
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0! (when checking argument for argument weight in method wrapper__cudnn_convolution)</p>

ptrblck: <p>Based on the stacktrace the error is raised in <a href="https://github.com/isl-org/DPT/blob/f43ef9e08d70a752195028a51be5e1aff227b913/dpt/vit.py#L175">this line of code</a> as it seems that <code> x = self.patch_embed.backbone(x)</code> creates the device mismatch internally.<br>
I’m not familiar with the <code>dpt</code> repository, so don’t know how these methods are used exactly.<br>
<code>nn.DataParallel</code> will chunk the input batch to the forward pass and will send each slice to the corresponding GPU which already holds the model clones.<br>
This would mean that you should call directly to the model and not into its <code>model.forward</code> method.<br>
Also, make sure that no new tensors in the <code>forward</code> are created on the default GPU. If newly created tensors are needed, then use the <code>.device</code> attribute or any parameter or the input tensor.</p>