Solved by iamneerav in post #8 
              
              
                So, I corrected the loss function. The error was in the script. Earlier, I was enabling the gradients inside each training batch as follows: 
for epoch in range(params.train.num_epochs):
    for train_batch, train_data in enumerate(trainloader):
        model.train()
        optimizer.zero_grad()
  …