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, 1, 1)
Codes is obtained as
codes = net.encoder(input) out = net.decoder(codes)
net is an encoder-decoder network
Not sure how to fix this error, any help is appreciated!