Yes, the sum or mean could should work.



 Lukens:

I have a tensor as large as my chosen batch (2), but I should have a single value to calculate the loss or am I wrong?


Since you are using reduction='none', an unreduced loss is expected for each criterion, so you could use the default mean reduction instead.