I am using Dataparallel module to train my network. However, I am facing the gpu device error when the MultivariateNormal distribution module is used. with my tensors.
# Using DataParallel
prior = MultivariateNormal(torch.zeros(dim).to(device), torch.eye(dim).to(device))
model = Flow(dim, prior, n_block)
model = model.to(device)
if use_cuda and torch.cuda.device_count()>1:
     model = DataParallel(model, device_ids=[0, 1, 2, 3])

# Calculating log probability of my Multivariate distribution
logprob = prior.log_prob(z).view(x.size(0), -1).sum(1)
# Error
File "/data/saandeepaath/flow_based/modules/flows.py", line 95, in forward
logprob = self.prior.log_prob(z).view(x.size(0), -1).sum(1)
File "/home/saandeepaath/.local/lib/python3.7/site-packages/torch/distributions/multivariate_normal.py", line 207, in log_prob
    diff = value - self.loc
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!