import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torchvision import datasets, transforms
 from torch.utils.data.dataloader import DataLoader

 import copy

 devices    = [0,1]
 batch_size = 64
 sim_type   = 'snn'
 timesteps  = 100

 device = 'cuda:' + str(devices[0])

 def test(model, epoch, test_loader):
     print('Testing model..')

     total = 0
     with torch.no_grad():
         model.eval()

         for i in range(100):
             # looping over dataset (start)
             for batch_idx, (data, target) in enumerate(test_loader):

                 if sim_type == 'snn':
                     data = data.repeat(timesteps, 1, 1, 1, 1)
                     data = data.swapaxes(0, 1)
                     B,T,C,H,W = data.size()
                 else:
                     B,C,H,W   = data.size()

                 data, target = data.to(device), target.to(device)
                 if batch_idx == 0:
                     print('data: {}'.format(data.size()))

                 output  = model(data)
                 loss    = F.cross_entropy(output, target)
                 pred    = output.max(1,keepdim=True)[1]
                 correct = pred.eq(target.data.view_as(pred)).cpu().sum()
                 total   += B
                 top1    = correct / total

                 print('Batch: {}, Accuracy: {}/{}({:.4f}), Loss: {:.4f}'
                     .format( batch_idx, correct.item(), B, top1, loss,))
             # looping over dataset (end)

 def load_cifar100():
     normalize = transforms.Normalize(mean = [0.5, 0.5, 0.5], std = [0.5, 0.5, 0.5])
     transform = transforms.Compose([transforms.ToTensor(), normalize])
     testset   = datasets.CIFAR100(root='./Datasets/cifar_data', \
             train=False, download=False, transform=transform)
     return testset

 class ANN(nn.Module):
     def __init__(self):
         super().__init__()

         self.ann_dict = nn.ModuleDict({
             'conv1': nn.Conv2d(3, 64, 3, padding=1, stride=1, bias=False),
             'relu1': nn.ReLU(inplace=True),
             'avg1' : nn.AvgPool2d(2, 2),

             'conv2': nn.Conv2d(64, 128, 3, padding=1, stride=1, bias=False),
             'relu2': nn.ReLU(inplace=True),
             'avg2' : nn.AvgPool2d(2, 2),

             'flat3': nn.Flatten(),
             'lin4' : nn.Linear(8192, 100, bias=False),
             })

     def forward(self, x):
         out_prev = x
         for k, v in self.ann_dict.items():
             out_prev = self.ann_dict[k](out_prev)
         return out_prev

 class SNN(nn.Module):
     def __init__(self, ann, device, in_sz):
         super(SNN, self).__init__()

         num_relus = 0
         for k, v in ann.named_modules():
             if isinstance(v, nn.ReLU):
                 num_relus += 1

         self.device   = device
         thresholds    = torch.ones(num_relus, device=device)
         self.snn_cell = SNNCell(ann, thresholds, device=device, in_sz=in_sz)

     def reset_vmems(self):
         self.snn_cell.reset_vmems()

     def get_vmems(self):
         return self.snn_cell.get_vmems()

     def forward(self, x):
         self.reset_vmems()
         vmems = self.get_vmems()
         return self.snn_cell(x, *vmems)

 class SNNCell(nn.Module):
     def __init__(self, ann, thresholds, device, in_sz):
         super(SNNCell, self).__init__()

         self.device   = device
         self.in_sz    = in_sz
         self.snn_dict = self.convert_to_snn(ann, thresholds)

     def convert_to_snn(self, ann, thr):
         snn_dict   = copy.deepcopy(ann.ann_dict)
         relu_idx   = 0
         in_sz      = self.in_sz
         for k, v in snn_dict.items():
             out_sz = v.cpu()(torch.zeros(in_sz)).size()
             if isinstance(v, nn.ReLU):
                 snn_dict[k] = SpikeRelu(thr[relu_idx], sz=in_sz, device=self.device,)
                 relu_idx += 1
             in_sz = out_sz
         return snn_dict

     def reset_vmems(self):
         for k, v in self.snn_dict.items():
             if isinstance(v, SpikeRelu):
                 self.snn_dict[k].reset_vmem()
     def get_vmems(self):
         vmems = []
         for k, v in self.snn_dict.items():
             if isinstance(v, SpikeRelu):
                 vmems.append(v.vmem.clone())
         return vmems

     def forward(self, x, *mem):
         T = x.size(1)

         L       = len(mem)
         mem_in  = list(mem[:])
         mem_out = [None] * L

         for t in range(T):
             spk_layer_idx  = 0
             outprev        = x[:,t]
             for name, layer in self.snn_dict.items():
                 if isinstance(layer, SpikeRelu):
                     outprev, mem_out[spk_layer_idx] = layer(outprev, mem_in[spk_layer_idx])
                     spk_layer_idx += 1
                 else:
                     outprev = layer(outprev)
             mem_in = mem_out[:]
         return outprev

 class SpikeRelu(nn.Module):
     def __init__(self, vth, sz, device='cuda:0'):
          super(SpikeRelu, self).__init__()

          self.threshold = vth
          self.sz        = sz
          self.device    = device
          self.reset_vmem()

     def reset_vmem(self):
         self.vmem = torch.zeros(*self.sz, device=self.device, requires_grad=True)

     def forward(self, x, mem_in):
         mem_thr = (mem_in / self.threshold) - 1.0
         op_spk  = torch.zeros_like(mem_thr)
         op_spk[mem_thr > 0] = 1.0
         rst     = self.threshold * (mem_thr > 0).float()
         mem_out = mem_in + x - rst
         return op_spk, mem_out

     def __repr__(self):
         return 'SpikeRelu(v_th : {:.3f}, vmem : {})'.format(self.threshold, torch.sum(self.vmem))


Hello everyone,
I am trying to apply torch.nn.DataParallel to a custom module (class SNN) whose simplified definition is provided above. I perform inference on this network using the following code:
 test_set    = load_cifar100()
 test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False)
 model = ANN()

 print('devices ', devices, 'device ', device)
 in_sz = (1,3,32,32)
 if sim_type == 'snn':
     model = SNN(model, device, in_sz, )

 model = torch.nn.DataParallel(model, device_ids=devices)
 model = model.to(device)

 print(model)
 test(model, 0, test_loader)

This keeps giving the following error: RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!
I believe the error is due to the way I am initializing my  vmem  in class SpikeRelu but I am not sure how to fix it. Please help!!