I’m not sure, what’s causing the issue, as the shapes work for me:
nb_classes = 10
output = torch.randn(1, nb_classes, 384,384,384, requires_grad=True)
target = torch.randint(0, nb_classes, (1, 384, 384, 384))
criterion = nn.CrossEntropyLoss()
loss = criterion(output, target)
loss.backward()

Could you double check the shapes of the output and target?