Solved by justusschock in post #6 
              
              
                Probably you might be able to achieve this with something simple like this: 
from torchvision.utils import save_image

# assuming your uncertainty matrix is stored in variable matrix 
save_image(1-matrix, "SPECIFY/YOUR/SAVE/PATH/uncertainty.png")

EDIT: Or if you want to plot it directly (e.g. by us…