Yorgos_Pantis:

So my output should be 6, since we have 6 classes.


Yes, this should be the case, but did you also verify it by checking the shape?
Based on the error message, this doesn’t seem to be the case as seen here:
criterion = nn.CrossEntropyLoss(weight=torch.randn(6))
x = torch.randn(1, 2)
target = torch.randint(0, 6, (1,))
out = criterion(x, target)
> RuntimeError: weight tensor should be defined either for all 2 classes or no classes but got weight tensor of shape: [6]