Thank you for your quick response. I included your suggestion and rerun the code. I am testing the model with an image I know that has 2 classes
import torchvision
import numpy as np
from torch.utils.data.sampler import SubsetRandomSampler
from PIL import Image
import matplotlib.pyplot as plt
import torchvision.transforms as T

def decode_segmap(image, nc=2):
    label_colors = np.array([(0, 0, 0),(255,255,255)])
    r = np.zeros_like(image).astype(np.uint8)
    g = np.zeros_like(image).astype(np.uint8)
    b = np.zeros_like(image).astype(np.uint8)

    for l in range(0, nc):
        idx = image == l
        r[idx] = label_colors[l, 0]
        g[idx] = label_colors[l, 1]
        b[idx] = label_colors[l, 2]

    rgb = np.stack([r, g, b], axis=2)
    return rgb
#import pic directly, webmethod check
import cv2
img = cv2.imread('/home/ubuntu/Downloads/Brain/test0716/train/slice_src_BN01002_032.png')
label = cv2.imread('/home/ubuntu/Downloads/Brain/test0716/trainannot/slice_src_BN01002_032.png')
plt.imshow(img); plt.show()
plt.imshow(label); plt.show()
trf = T.Compose([T.ToTensor()])
inp = trf(img).unsqueeze(0).cuda()
fcn = torch.load('/home/ubuntu/Downloads/Brain/test0716/mod.pth')
fcn.eval()
sam_out = fcn(inp)['out']
om = torch.argmax(sam_out.squeeze(), dim=0).cpu().numpy()
rgb = decode_segmap(om)
print(rgb.shape)
print (np.unique(rgb))
plt.imshow(rgb);


This is the code I use to visualize the result, The prediction gives me an image with only 1 label. I greatly appreciate your help. Please let me know if you have any suggestions.
Thanks
Nishanth