Thanks a lot. I did the following and obtaining a black image only so I guess my model sucks and accuracy is really not the right metric to judge unbalanced datasets.
def predict_fundus_image(img_path,model):
    img = Image.open(img_path)
    
    transform = transforms.Compose([transforms.Resize((512,512)),
                                    transforms.ToTensor()])
    img_transformed = transform(img)
    img_transformed = transform(img).unsqueeze_(0)
    if train_on_gpu:
        img_transformed = img_transformed.cuda()
    if train_on_gpu:
        model = model.cuda()
    output = model(img_transformed)
    output = torch.argmax(output,1)
    return output

img = predict_fundus_image(img_path,model)
img = img.squeeze(0)
mapping = {(0, 0, 0): 0, (0, 0, 255): 1, (255, 0, 0): 2, (255, 255, 255): 3}
rev_mapping = {mapping[k]: k for k in mapping}
pred = img
pred_image = torch.zeros(3, img.size(0), img.size(1), dtype=torch.uint8)
for k in rev_mapping:
    print(k)
    pred_image[:, pred==k] = torch.tensor(rev_mapping[k]).byte().view(3, 1)
plt.imshow(pred_image.permute(1, 2, 0).numpy())