Solved by ptrblck in post #8 
              
              
                I’m not sure, why it’s not working. 
If you would like to visualize both probability maps for the two classes, your code should work: 
plt.figure()
plt.imshow(torch.exp(outputs[0,0,:,:]).detach().cpu())  # plot class0
plt.figure()
plt.imshow(torch.exp(outputs[0,1,:,:]).detach().cpu())  # plot class1…