The following code seems to be working  but ground truth and prediction mask images are condensed 
Screen Shot 2020-03-06 at 4.26.05 PM1684×358 102 KB
 . So weird.
for i in range(2):
    n = np.random.choice(len(test_dataset))
    image_vis = test_dataset_vis[n][0].astype('uint8')
    
    image, gt_mask = test_dataset[n]
    
    gt_mask1 = gt_mask[...,0].squeeze()
    gt_mask2 = gt_mask[...,1].squeeze()
    gt_mask3 = gt_mask[...,2].squeeze()
    gt_mask4 = gt_mask[...,3].squeeze()

    x_tensor = torch.from_numpy(image).to(DEVICE).unsqueeze(0)
    pr_mask = model.predict(x_tensor)
    
    pr_mask1 = (pr_mask[...,0].squeeze().cpu().numpy().round())
    pr_mask2 = (pr_mask[...,1].squeeze().cpu().numpy().round())
    pr_mask3 = (pr_mask[...,2].squeeze().cpu().numpy().round())
    pr_mask4 = (pr_mask[...,3].squeeze().cpu().numpy().round())

        
    visualize(
        image=image_vis,
        gt_mask1 = gt_mask1,
        gt_mask2 = gt_mask2,
        gt_mask3 = gt_mask3,
        gt_mask4 = gt_mask4,
        pr_mask1 = pr_mask1,
        pr_mask2 = pr_mask2,
        pr_mask3 = pr_mask3,
        pr_mask4 = pr_mask4
    )

# helper function for data visualization
def visualize(**images):
    """PLot images in one row."""
    n = len(images)
    plt.figure(figsize=(16, 5))
    for i, (name, image) in enumerate(images.items()):
        plt.subplot(1, n, i + 1)
        plt.xticks([])
        plt.yticks([])
        plt.title(' '.join(name.split('_')).title())
        plt.imshow(image)
    plt.show()