Hi Richie!



 richieYT-wan:

Is there a way to use the CrossEntropyLoss for a 2D target (so final dimension would be 3D, (batchsize, dim1, dim2))


Yes.  CrossEntropyLoss 3 supports what it calls the “K-dimensional case.”
Note, pytorch’s CrossEntropyLoss does not accept a one-hot-encoded
target – you have to use integer class labels instead.
Let’s call your value 23 length.  Your input (the prediction generated
by your network) should have shape
[mini_batch, N_classes, length]
Your target (the ground-truth labels) should have shape
[mini_batch, length] without an N_classes dimension and should
be integer class labels in [0, N_classes - 1], inclusive.  (In particular,
they are not one-hot encoded.)
Last, do not pass the output of your network through softmax();
CrossEntropyLoss has, in effect, softmax() built in (and expects
its input to be raw-score logits).
Best.
K. Frank