nn.DataParallel multigpu: "...found two devices cuda:0 and cuda:1!"

I have a model net that I am trying to train on multiple GPUs. I used DataParallel to enforce this, however during training when I run a forward pass through this model I am getting the following error

RuntimeError: Expected all tensors to be on the same device, 
but found at least two devices, cuda:1 and cuda:0!

the piece of code that is causing this error is

 codes = codes + self.latent_avg.repeat(codes.shape[0], 1, 1)

Codes is obtained as

codes = net.encoder(input)
out = net.decoder(codes)

where net is an encoder-decoder network

Not sure how to fix this error, any help is appreciated!

Can you check if latent_avg is on a different device than codes? To run DP, net needs to be on cuda:0 => codes would be on cuda:0 so presumably self.latent_avg is on cuda:1 for you to get this error.

if self.latent_avg.device.index != 0:
  self.latent_avg = self.latent_avg.cuda(0)

codes = codes + self.latent_avg.repeat(codes.shape[0], 1, 1)