I have the same problem with this code: (pytorch from source)
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
class MyModule(nn.Module):
def __init__(self):
super(MyModule, self).__init__()
self.conv = nn.Conv2d(3, 3, 3, 1, 0)
def forward(self, x):
size = [int(s * 0.5) for s in x.shape[2:]]
a = self.conv(x)
b = F.upsample(x, size=size, mode='bilinear', align_corners=True)
b = self.conv(b)
c = F.upsample(b, size=a.shape[2:], mode='bilinear', align_corners=True)
return a, b, c
data = torch.rand(5, 3, 32, 32).cuda()
data = Variable(data)
model = MyModule()
model = nn.DataParallel(model)
model.cuda()
outputs = model(data)
loss = 0
target_a = np.random.randint(0, 3, size=(5, 30, 30))
target_a = torch.from_numpy(target_a).long().cuda()
target_a = Variable(target_a, requires_grad=False)
loss += F.nll_loss(F.log_softmax(outputs[0], dim=1), target_a, ignore_index=-1)
target_b = np.random.randint(0, 3, size=(5, 14, 14))
target_b = torch.from_numpy(target_b).long().cuda()
target_b = Variable(target_b, requires_grad=False)
loss += F.nll_loss(F.log_softmax(outputs[1], dim=1), target_b, ignore_index=-1)
target_c = np.random.randint(0, 3, size=(5, 30, 30))
target_c = torch.from_numpy(target_c).long().cuda()
target_c = Variable(target_c, requires_grad=False)
loss += F.nll_loss(F.log_softmax(outputs[2], dim=1), target_c, ignore_index=-1)
loss.backward()
Output:
Traceback (most recent call last):
File "test.py", line 68, in <module>
loss.backward()
File "/opt/conda/envs/pytorch-py36/lib/python3.6/site-packages/torch/tensor.py", line 93, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph)
File "/opt/conda/envs/pytorch-py36/lib/python3.6/site-packages/torch/autograd/__init__.py", line 89, in backward
allow_unreachable=True) # allow_unreachable flag
RuntimeError: arguments are located on different GPUs at /root/pytorch/aten/src/THC/generated/../generic/THCTensorMathPointwise.cu:233
Process finished with exit code 1
Any idea? @ptrblck