Cross entropy loss for 3D tensor

I have an output tensor (both target and predicted) of dimension (32 x 8 x 5000). Here, the batch size is 32, the number of classes is 5000 and the number of points per batch is 8. I want to calculate CELoss on this in such a way that, the loss is computed for every point and then averaged across 8 of them. How can I do this? Let me know if my question isn’t clear or ambiguous

CrossEntropyLoss expects an input of dim = (N, C) where N=batch_size and C=no. of classes,
and a target of dim = (N,).

Additional dimensions are used for “K-dimensional loss” as stated in the docs.

It’s not quite clear to me what the second dimension (8) represents but potentially check out the reduction parameter, that, by default is “mean” so you shouldn’t need modification for your use.

1 Like