Pytorch crossentropy loss with 3d input
I have a network which outputs a 3D tensor of size (batch_size, max_len, num_classes). My groud truth is in the shape (batch_size, max_len). If I do perform one-hot encoding on the labels, it'll be of shape (batch_size, max_len, num_classes) i.e the values in max_len are integers in the range [0, num_classes]. Since the original code is too long, I have written a simpler version that reproduces the original error.
criterion = nn.CrossEntropyLoss()
batch_size = 32
max_len = 350
num_classes = 1000
pred = torch.randn([batch_size, max_len, num_classes])
label = torch.randint(0, num_classes,[batch_size, max_len])
pred = nn.Softmax(dim = 2)(pred)
criterion(pred, label)

the shape of pred and label are respectively,torch.Size([32, 350, 1000]) and torch.Size([32, 350])
The error encountered is
ValueError: Expected target size (32, 1000), got torch.Size([32, 350, 1000])

If I one-hot encode labels for computing the loss
x = nn.functional.one_hot(label)
criterion(pred, x)

it'll throw the following error
ValueError: Expected target size (32, 1000), got torch.Size([32, 350, 1000])


From the Pytorch documentation, CrossEntropyLoss expects the shape of its input to be (N, C, ...), so the second dimension is always the number of classes. Your code should work if you reshape preds to be of size (batch_size, num_classes, max_len).
