Solved by ptrblck in post #13 
              
              
                Ah ok, thanks for the info. 
It looks like a standard segmentation task. 
I would suggest to use nn.CrossEntropyLoss for your use case. 
Have a look at the following code snippet: 
n_class = 10
preds = torch.randn(4, n_class, 24, 24)
labels = torch.empty(4, 24, 24, dtype=torch.long).random_(n_class)…