Why is there a special case for 2D in nn.Softmax with dim=None?


I have a question about this code:

if ndim == 0 or ndim == 1 or ndim == 3:
    ret = 0
    ret = 1

Why are we picking 1 as a dimension for the 2D case?

From the Slack cross-post:

I would guess because 2D tensors often represent [batch_size, nb_classes] where dim1 contains the logits/features/etc.

1 Like

Thank you, that is helpful!