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 shapes.
Let me know, if it would work for you.