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