Hi @ailzhang, thanks for the reply. Here is my model code :
#!/usr/bin/env python

import torch
import torch.nn as nn
from torch.autograd import Variable

import numpy as np

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class VAE(nn.Module):
    def __init__(self, nc, ngf, ndf, latent_variable_size):
        super(VAE, self).__init__()

        self.nc = nc
        self.ngf = ngf
        self.ndf = ndf
        self.latent_variable_size = latent_variable_size

        # encoder
        self.e1 = nn.Conv2d(nc, ndf, 4, 2, 1)
        self.bn1 = nn.BatchNorm2d(ndf)

        self.e2 = nn.Conv2d(ndf, ndf*2, 4, 2, 1)
        self.bn2 = nn.BatchNorm2d(ndf*2)

        self.e3 = nn.Conv2d(ndf*2, ndf*4, 4, 2, 1)
        self.bn3 = nn.BatchNorm2d(ndf*4)

        self.e4 = nn.Conv2d(ndf*4, ndf*8, 4, 2, 1)
        self.bn4 = nn.BatchNorm2d(ndf*8)

        self.e5 = nn.Conv2d(ndf*8, ndf*8, 4, 2, 1)
        self.bn5 = nn.BatchNorm2d(ndf*8)

        self.fc1 = nn.Linear(ndf*8*4*4, latent_variable_size)
        self.fc2 = nn.Linear(ndf*8*4*4, latent_variable_size)
        self.fc3 = nn.Linear(ndf*8*4*4, latent_variable_size)

        # decoder
        self.d1 = nn.Linear(latent_variable_size, ngf*8*2*4*4)

        self.up1 = nn.UpsamplingNearest2d(scale_factor=2)
        self.pd1 = nn.ReplicationPad2d(1)
        self.d2 = nn.Conv2d(ngf*8*2, ngf*8, 3, 1)
        self.bn6 = nn.BatchNorm2d(ngf*8, 1.e-3)

        self.up2 = nn.UpsamplingNearest2d(scale_factor=2)
        self.pd2 = nn.ReplicationPad2d(1)
        self.d3 = nn.Conv2d(ngf*8, ngf*4, 3, 1)
        self.bn7 = nn.BatchNorm2d(ngf*4, 1.e-3)

        self.up3 = nn.UpsamplingNearest2d(scale_factor=2)
        self.pd3 = nn.ReplicationPad2d(1)
        self.d4 = nn.Conv2d(ngf*4, ngf*2, 3, 1)
        self.bn8 = nn.BatchNorm2d(ngf*2, 1.e-3)

        self.up4 = nn.UpsamplingNearest2d(scale_factor=2)
        self.pd4 = nn.ReplicationPad2d(1)
        self.d5 = nn.Conv2d(ngf*2, ngf, 3, 1)
        self.bn9 = nn.BatchNorm2d(ngf, 1.e-3)

        self.up5 = nn.UpsamplingNearest2d(scale_factor=2)
        self.pd5 = nn.ReplicationPad2d(1)
        self.d6 = nn.Conv2d(ngf, nc, 3, 1)

        self.leakyrelu = nn.LeakyReLU(0.2)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def encode(self, x):
        h1 = self.leakyrelu(self.bn1(self.e1(x)))
        h2 = self.leakyrelu(self.bn2(self.e2(h1)))
        h3 = self.leakyrelu(self.bn3(self.e3(h2)))
        h4 = self.leakyrelu(self.bn4(self.e4(h3)))
        h5 = self.leakyrelu(self.bn5(self.e5(h4)))
        h5 = h5.view(-1, self.ndf*8*4*4)

        return self.fc1(h5), self.fc2(h5), self.fc3(h5)

    def reparameterize(self, mu, logvar):
        # std = logvar.mul(0.5).exp_()
        # eps = torch.cuda.FloatTensor(std.size()).normal_()
        # eps = Variable(eps).cuda()
        # return eps.mul(std).add_(mu)
        if self.training:
            std = torch.exp(0.5*logvar)
            eps = torch.randn_like(std)
            return eps.mul(std).add_(mu)
        else:
            return mu

    def decode(self, common_factor, varying_factor):
        h1_c = self.relu(self.d1(common_factor))
        h1_v = self.relu(self.d1(varying_factor))
        # h1 = torch.cat([h1_c, h1_v], dim=0)
        h1 = h1_c + h1_v
        h1 = h1.view(-1, self.ngf*8*2, 4, 4)
        h2 = self.leakyrelu(self.bn6(self.d2(self.pd1(self.up1(h1)))))
        h3 = self.leakyrelu(self.bn7(self.d3(self.pd2(self.up2(h2)))))
        h4 = self.leakyrelu(self.bn8(self.d4(self.pd3(self.up3(h3)))))
        h5 = self.leakyrelu(self.bn9(self.d5(self.pd4(self.up4(h4)))))

        return self.sigmoid(self.d6(self.pd5(self.up5(h5))))

    def get_latent_var(self, x):
        common_factor, mu, logvar = self.encode(x.view(-1, self.nc, self.ndf, self.ngf))
        z = self.reparametrize(mu, logvar)
        return z

    def forward(self, x):
        common_factor, mu, logvar = self.encode(x.view(-1, self.nc, self.ndf, self.ngf))
        z = self.reparametrize(mu, logvar)
        res = self.decode(z)
        return res, mu, logvar



model = VAE(nc=3, ngf=128, ndf=128, latent_variable_size=500)
model = nn.DataParallel(model).to(device)

model.train()

train_data_1 = np.random.randn(10, 20, 3, 128, 128)
train_data_2 = np.random.randn(10, 20, 3, 128, 128)

train_data_1 = torch.from_numpy(train_data_1).float()
train_data_2 = torch.from_numpy(train_data_2).float()

# train_data_1 = torch.unsqueeze(train_data_1, 0)
# train_data_2 = torch.unsqueeze(train_data_2, 0)

def vae_loss_function(image1, decoded_1, image2, decoded_2, mu_1, logvar_1):
	MSE = nn.MSELoss()
	mse_loss_1 = MSE(image1, decoded_1.detach())
	mse_loss_2 = MSE(image2, decoded_2.detach())
	# print("Shape of image is {} and decoded tensor is {}".format(image1.shape, decoded_1.shape))
	# mse_loss_1 = torch.sum((image1 - decoded_1)**2) / image1.data.nelement()
	# mse_loss_2 = torch.sum((image2 - decoded_2)**2) / image2.data.nelement()        
	kl_div = torch.sum(-0.5 * torch.sum(1 + 2 * logvar_1 - mu_1**2 - torch.exp(2*logvar_1),1))

	total_loss = mse_loss_1 + mse_loss_2 + kl_div
	return total_loss

num_epochs = 100
lr = 0.001

optimizer = torch.optim.Adam(model.parameters(), lr)

train_data_1 = train_data_1.to(device)
train_data_2 = train_data_2.to(device)

for epoch in range(num_epochs):
	for batch_idx, image in enumerate(train_data_1):
		# print('hello batch {}'.format(batch_idx))
		common_factor_1, mu_1, logvar_1 = model.module.encode(image)
		reparam_out_1 = model.module.reparameterize(mu_1, logvar_1)
		common_factor_2, mu_2, logvar_2 = model.module.encode(train_data_2[batch_idx])
		reparam_out_2 = model.module.reparameterize(mu_2, logvar_2)

		decoded_1 = model.module.decode(common_factor_2, reparam_out_1)
		decoded_2 = model.module.decode(common_factor_1, reparam_out_2)

		loss = vae_loss_function(image, decoded_1, train_data_2[batch_idx], decoded_2, mu_1, logvar_1)
		
		optimizer.zero_grad()
		loss.backward()
		train_loss = loss.item()
		optimizer.step()

		if batch_idx % 10 == 0:
			print(
				'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tlr: {:.6f}'.format(
				epoch+1, batch_idx * len(train_data_1), len(train_data_1),
				100. * batch_idx / len(train_data_1),
				loss.item() / len(train_data_1), lr)
			)

This code is not able to use more than one gpu available to it. What is wrong in here?