Solved by KFrank in post #2 
              
              
                Hi Granth! 

The short answer is that you are calling python’s max() function, 
rather than pytorch’s torch.max() tensor function.  This is causing 
you to calculate softmax() for a tensor that is all zeros. 
You have two issues: 
First is the use of pytorch’s max(). max() doesn’t understand 
tensor…