I’m not sure, what Input and output are.
Based on the shapes it seems that Input is actually your model prediction, while output seems to be the target. Is this correct?
Here is a small example of using nn.CrossEntropyLoss for images, e.g. in a segmentation use case:
batch_size = 1
c, h, w = 1, 10, 10
nb_classes = 3
x = torch.randn(batch_size, c, h, w)
target = torch.empty(batch_size, h, w, dtype=torch.long).random_(nb_classes)

model = nn.Conv2d(c, nb_classes, 3, 1, 1)
criterion = nn.CrossEntropyLoss()

output = model(x)
loss = criterion(output, target)
loss.backward()

Let me know, if your use case differs.