Solved by ptrblck in post #6 
              
              
                I think nn.MultiMarginLoss would be the suitable criterion: 

Creates a criterion that optimizes a multi-class classification hinge loss (margin-based loss) between input x (a 2D mini-batch Tensor) and output y 

Based on the shape information it should also work for your current output and target s…