Solved by KFrank in post #2 
              
              
                Hi Abishek! 

I assume that criterion_seg is something you wrote yourself 
that basically forwards its arguments to F.cross_entropy(). 
You most likely should pass Model Output in place of x as 
the first argument of criterion_seg, and don’t use argmax() 
at all. 

This is (most likely) telling you …