Multi-Class cross entropy loss batch size error

Hi,
I am working on nuscenes dataset and for one of the output head using cross entropy loss.
input size torch.Size([8, 3, 10, 159, 159])
target size torch.Size([8, 10, 159, 159])
8 - batch size
3 - classes (specific to head)
10 - d1 ( these are overall classes; for each class, we can have 3 values specifically as mentioned above)
159 - d2 (height)
159 -d3 (width)
We have output heads in our model and the model predicts all the attributes for the main 10 different classes.

If I am using a batch size of 1 then everything is working fine.
input [1, 3, 10 159, 159]
target [1, 10, 159, 159]

But when I increase the batch size cross-entropy loss results in the error "RuntimeError: The size of tensor a (8) must match the size of tensor b (10) at non-singleton dimension 1"

I am using torch 1.7.1.

Am I missing something from the definition of cross-entropy loss?
Any suggestion on how I can correct it?

Both approaches should work as seen here:

criterion = nn.CrossEntropyLoss()

input = torch.randn(1, 3, 10, 159, 159, requires_grad=True)
target = torch.randint(0, 3, (1, 10, 159, 159))
loss = criterion(input, target)
print(loss)
# tensor(1.3898, grad_fn=<NllLoss2DBackward0>)

input = torch.randn(8, 3, 10, 159, 159, requires_grad=True)
target = torch.randint(0, 3, (8, 10, 159, 159))
loss = criterion(input, target)
print(loss)
# tensor(1.3896, grad_fn=<NllLoss2DBackward0>)

Are you changing the shape of these tensors internally somehow?