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])