I still couldn’t figure out how can I visualise the output of network after training. This code will visualize the raw output but I don’t know how can I display all dim of image, at the moment will display only one channel in plt.imshow(outputs[0,0,:,:].detach().cpu()) while the shape is #print(outputs.shape) # torch.Size([1, 2, 240, 320]) it is the same issue with plt.imshow(t_image[0,0,:,:].detach().cpu()) while the shape is #print(t_image.size()) # torch.Size([1, 1, 240, 320])
model.eval()
    total = 0
    test_loss = 0
    correct = 0
    count = 0
    #iterate through test dataset
    for vi, data in enumerate(test_loader):
                
                t_image, mask = data
                #print(t_image.shape) # torch.Size([1, 1, 240, 320])
                t_image, mask = t_image.to(device), mask.to(device)
                with torch.no_grad():

                    outputs = model(t_image)
                    #print(outputs.shape) # torch.Size([1, 2, 240, 320])
                    test_loss += criterion(outputs, mask).item() / len(test_loader)

                    pr = torch.exp(outputs) # get the exp of output and will give the probability map of outputs
                    # The outputs are energies for the 2 classes.
                    # Higher the energy for a class, the more the network thinks that the image is of the particular class. So, let’s get the index of the highest energy:     
                    _, predicted = torch.max(outputs.data, 1)  
                    
                    total += mask.size(0)
                    correct += predicted.eq(mask.data).sum().item()
                    accuracy = 100 * correct / total
                    predict = predicted.eq(mask.data)
                    #print(predict.shape) # torch.Size([1, 240, 320])
                    count +=1                
                    print(count, "Test Loss: {:.3f}".format(test_loss), "Test Accuracy: %d %%" % (accuracy))

    
    plt.figure()
    plt.subplot(1, 4, 1)
    #print(outputs.shape) # torch.Size([1, 2, 240, 320])
    plt.imshow(outputs[0,0,:,:].detach().cpu())
    plt.title('DL raw output')
    plt.subplot(1, 4, 2)
    plt.imshow(predict.detach().cpu().squeeze())
    plt.title('DL prediction')
    plt.subplot(1, 4, 3)
    plt.imshow(mask.detach().cpu().squeeze())
    plt.title('ground truth')
    plt.subplot(1, 4, 4)
    #print(t_image.size()) # torch.Size([1, 1, 240, 320])
    plt.imshow(t_image[0,0,:,:].detach().cpu())
    plt.title('original input image')

@ptrblck Could you please help to sort out this issue? I did try permute(1, 2, 0)) but couldn’t solve the issue