You don’t really need to view the tensor in a different way. Cross entropy is now equipped to handle multiple dimensions. The code below demonstrates its working.
import torch
import torch.nn as nn
criterion = nn.CrossEntropyLoss()
batch_size = 16
no_of_classes = 5
input = torch.randn(batch_size,no_of_classes,31)
target = torch.randint(0,4,(batch_size,31))
criterion(input,target)