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!