I get
RuntimeError: arguments are located on different GPUs at /pytorch/torch/lib/THC/generic/THCTensorIndex.cu:390
When I use index_selct function with torch.nn.DataParallel(net).
During forwarding pass of my model. I use tenspr.index_select() to flatten the output feature map of a Conv layer to feed it into the linear layer. Basically, I am implementing next CONV layer with a linear. Which I am going to use to increase the flexibility of my model. It works fine on one GPU and able to reproduce results of next CONV layer using index_select and the linear layer on top. Here is code part where I get this error.
import torch.nn as nn
import torch, math
from torch.autograd import Variable
class SpatialPool(nn.Module):
def __init__(self, amd0=225, kd=3):
super(SpatialPool, self).__init__()
print('*** spatial_pooling.py : __init__() ***', amd0,kd,fd)
self.use_gpu = True
self.amd0 = amd0 #225
self.kd = kd
self.padding = nn.ReplicationPad2d(1).cuda()
ww = hh = int(math.sqrt(amd0)) ## 15
counts = torch.LongTensor(amd0,kd*kd) ## size [225,9]
v = [[(hh+2)*i + j for j in range(ww+2)] for i in range(hh+2)]
count = 0
for h in range(1,hh+1):
for w in range(1,ww+1):
counts[count,:] = torch.LongTensor([v[h - 1][w - 1], v[h - 1][w], v[h - 1][w + 1],
v[h][w - 1], v[h][w], v[h][w + 1],
v[h + 1][w - 1], v[h + 1][w], v[h + 1][w + 1]])
count += 1
self.counts = counts.cuda()
def forward(self, fm):
fm = self.padding(fm) ## FM is Variable of size[batch_size,512,15,15]
fm = fm.permute(0, 2, 3, 1).contiguous()
fm = fm.view(fm.size(0), -1, fm.size(3))
print('fm size and max ', fm.size(), torch.max(self.counts))
pfm = fm.index_select(1,Variable(self.counts[:,0]))
for h in range(1,self.kd*self.kd):
pfm = torch.cat((pfm,fm.index_select(1, Variable(self.counts[:, h]))),2)
# print('pfm size:::::::: ', pfm.size()) #[batch_size,225,512*9]
return pfm
Here fm is a matrix of size.
Full error message looks like below
File "/home/gurkirt/Dropbox/sandbox/ssd-pytorch-linear/layers/modules/feat_pooling.py", line 48, in forward
pfm = self.spatial_pool1(fm) # pooled feature map
File "/usr/local/lib/python3.5/dist-packages/torch/nn/modules/module.py", line 224, in __call__
result = self.forward(*input, **kwargs)
File "/home/gurkirt/Dropbox/sandbox/ssd-pytorch-linear/layers/modules/spatial_pooling.py", line 36, in forward
pfm = fm.index_select(1,Variable(self.counts[:,0]))
File "/usr/local/lib/python3.5/dist-packages/torch/autograd/variable.py", line 681, in index_select
return IndexSelect.apply(self, dim, index)
File "/usr/local/lib/python3.5/dist-packages/torch/autograd/_functions/tensor.py", line 297, in forward
return tensor.index_select(dim, index)
RuntimeError: arguments are located on different GPUs at /pytorch/torch/lib/THC/generic/THCTensorIndex.cu:390