Solved by ptrblck in post #9 
              
              
                To calculate the confusion matrix you need the class predictions. Currently it looks like pred contains the logits or probabilities for two classes. 
Try to call torch.argmax(pred, 1) to get the predicted classes. 
Here is a small example: 
output = torch.randn(1, 2, 4, 4)
pred = torch.argmax(output…