Solved by mailcorahul in post #2 
              
              
                @alie There are two mistakes here. 

You apply softmax twice - once before calling your custom loss function and inside it as well.
You are not applying log to softmax output. Inside log_softmax(): 
return torch.log(torch.exp(x) / torch.sum(torch.exp(x), dim=1, keepdim=True))