@ptrblck one more question here, if I visualize the probability of output same as your example
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

it return a figure which is different with the figure that visualize by following example
    output = model(t_image)
    prob = torch.exp(output)
    prob_imgs = make_grid(prob.permute(1, 0, 2, 3))
    prob_imgs = prob_imgs.detach()
    plt.imshow(prob_imgs.permute(1, 2, 0))

what is cause to be different in these example, is it because of make_grid? because permute is just swap the first (batch size) and second dimensions (number of classes).
this image is from second example, for the first example the zeros value in figures will be 1 and 1 converted to zeros. (black and white will swap). which way is correct way to do this? thank you in advance