Frank,
Thanks for the detailed reply.



 KFrank:

Note, for the “K-dimensional loss” case, you will want to reshape your
logits to have shape [nBatch, 5, 6] (not [nBatch, 6, 5]).


I am somehow finding this order less intuitive for me - nClass appearing before nTask .
I will be using torch.nn.functional.softmax to use the logits outside the loss function. I will use the dim parameter to go along the correct dimension.
Thanks again.