RuntimeError                              Traceback (most recent call last)
<ipython-input-3-0c9f361b4bf0> in <module>
    260 if __name__ == "__main__":
    261     for epoch in range(1, epochs + 1):
--> 262         train(epoch)

<ipython-input-3-0c9f361b4bf0> in train(epoch)
    239 
    240 
--> 241         loss.backward(create_graph=False,retain_graph=True)
    242         optimizer.step()
    243 

~\Anaconda3\lib\site-packages\torch\tensor.py in backward(self, gradient, retain_graph, create_graph)
    193                 products. Defaults to ``False``.
    194         """
--> 195         torch.autograd.backward(self, gradient, retain_graph, create_graph)
    196 
    197     def register_hook(self, hook):

~\Anaconda3\lib\site-packages\torch\autograd\__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
     97     Variable._execution_engine.run_backward(
     98         tensors, grad_tensors, retain_graph, create_graph,
---> 99         allow_unreachable=True)  # allow_unreachable flag
    100 
    101 

RuntimeError: Function 'PowBackward0' returned nan values in its 0th output.

I have tried it on a sample data of 10 datapoints with parameter sharing,it works on that. but when I do increase the datasize with the parameter sharing this error occurs.
#model class
class Framework(nn.Module):
    def __init__(self,input_shape, representation_size, output_shape, fc1, fc21,fc22,fc3,fc4):
        super(Framework, self).__init__()
        self.fc1 = fc1
        self.fc21 = fc21
        self.fc22 = fc22
        self.fc3 = fc3
        self.fc4 = fc4
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        self.elu=nn.ELU()
    
    def encoder(self, x):
        """Encode a batch of samples, and return posterior parameters for each point."""
        h1 = self.relu(self.fc1(x))
        return self.fc21(h1), self.elu(self.fc22(h1))+1
    
    def decoder(self, z):
        h3 = self.relu(self.fc3(z))
        return self.sigmoid(self.fc4(h3))
        
    def reparam(self, mu, logvar):
        """Reparameterisation trick to sample z values. 
        This is stochastic during training,  and returns the mode during evaluation."""
        
        if self.training:
            std = logvar.mul(0.5).exp_()
            eps = Variable(std.data.new(std.size()).normal_())
            return eps.mul(std).add_(mu)
        else:
            return mu
        
    def get_z(self, x):
        """Encode a batch of data points, x, into their z representations."""
        
        mu, logvar = self.encoder(x.view(-1, input_shape))
        return self.reparam(mu, logvar)
    
    def forward(self, x):
        """Takes a batch of samples, encodes them, and then decodes them again to compare."""
        mu, logvar = self.encoder(x.view(-1, input_shape))
        z = self.reparam(mu, logvar)
        return self.decoder(z), mu, logvar,z

#parameter sharing
fc1 = nn.Linear(input_shape, 512,bias=True)
torch.nn.init.xavier_uniform_(fc1.weight)
fc21 = nn.Linear(512, representation_size,bias=True)
torch.nn.init.xavier_uniform_(fc21.weight)
fc22 = nn.Linear(512, representation_size,bias=True)
torch.nn.init.xavier_uniform_(fc22.weight)
fc3 = nn.Linear(representation_size, 512,bias=True)
torch.nn.init.xavier_uniform_(fc3.weight)
fc4 = nn.Linear(512, output_shape,bias=True)
torch.nn.init.xavier_uniform_(fc4.weight)

#Complete model
model1 = Framework(input_shape, representation_size, output_shape, fc1, fc21,fc22,fc3,fc4).to(device)
model2 = Framework(input_shape, representation_size, output_shape, fc1, fc21,fc22,fc3,fc4).to(device)
model3 = Framework(input_shape, representation_size, output_shape, fc1, fc21,fc22,fc3,fc4).to(device)

#deine optimizers

f_params=model1.parameters()
s_params=model2.parameters()
t_params=model3.parameters()

dvne_params=itertools.chain(f_params,s_params,t_params)

optimizer = optim.RMSprop(dvne_params, lr=learning_rate)

#train Model
epoch_loss=[]
Auc_score=[]
def train(epoch):
    
    #for param_group in optimizer.param_groups:
        ##print(param_group['lr'], "learning rate for Auto-Encoder.")
    model1.train()
    model2.train()
    model3.train()
    
    torch.autograd.set_detect_anomaly(True)
    train_loss,Auc=0,0
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    for batch_idx, data in enumerate(train_loader):
        print(batch_idx)
        #torch.autograd.set_detect_anomaly(False)
        data1 = data[0].to(device)
        data2 = data[1].to(device)
        data3 = data[2].to(device)

        data1 = Variable(data1)
        data2 = Variable(data2)
        data3 = Variable(data3)
        
        recon_batch1, mu1, logvar1,z1 = model1(data1.float())
        recon_batch2, mu2, logvar2,z2 = model2(data2.float())
        recon_batch3, mu3, logvar3,z3 = model3(data3.float())
        
        loss = wasserstein_loss(data1,recon_batch1,z1, mu1, logvar1, data2,recon_batch1,z2, mu2, logvar2, data3,recon_batch1,z3, mu3, logvar3)

        auroc=wasserstein_acc(data1,recon_batch1,data2,recon_batch2,data3,recon_batch3)
        auroc=100*auroc

        optimizer.zero_grad()

        loss.backward(retain_graph=True)
        optimizer.step()
        #print(recon_batch1, data1)
        train_loss+=loss.item()
        Auc+=auroc.item()
        if (batch_idx) % log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6}\tauc: {:.6}'.format(
                epoch, batch_idx*batch_size* len(data)/3, len(train_loader.dataset),
                100. * batch_idx / len(train_loader),
                loss.item()/ len(data),auroc))
            
    epochloss=(train_loss / (len(train_loader.dataset)))
    AUC=(Auc*batch_size/(len(train_loader.dataset)))
    #print('====> Epoch: {} Average loss: {:.6}\t Auc: {:.6f}'.format(epoch,epochloss,AUC))
    epoch_loss.append(epochloss)
    Auc_score.append(AUC)
    print('max_auc: {:.6}\tmin_loss: {:.6}\tmax_auc_epoch: {}\tmin_loss_epoch: {}'.format(max(Auc_score),min(epoch_loss),
                                                                Auc_score.index(max(Auc_score))+1,epoch_loss.index(min(epoch_loss))+1))
if __name__ == "__main__":
    for epoch in range(1, epochs + 1):
        train(epoch)
    epoch=np.arange(epochs)
    fig, ax = plt.subplots( nrows=1, ncols=1 )
    plt.plot(epoch,Auc_score, marker='o',color='r', label='Square')
    plt.plot(epoch,epoch_loss , marker='o',color='b', label='Square') 
    plt.xticks(epoch_loss, Auc_score)
    fig.savefig('curve.pdf')
    plt.close(fig)
    max(Auc_score),min(epoch_loss)