I have trained the model for multi-class image segmentation. My model is returning the torch tensor of shape torch.Size([4, 512, 512]).  Now, I want to convert this tensor to the image mask, as I want to visualise the output my model is predicting but I am unable to do it
Following is my image_to_class function which converts mask image to class labels-
def mask_to_class(self,mask):
    target = torch.from_numpy(mask)
    h,w = target.shape[0],target.shape[1]
    masks = torch.empty(h, w, dtype=torch.long)
    colors = torch.unique(target.view(-1,target.size(2)),dim=0).numpy()
    target = target.permute(2, 0, 1).contiguous()
    mapping = {tuple(c): t for c, t in zip(colors.tolist(), range(len(colors)))}
    for k in mapping:
        idx = (target==torch.tensor(k, dtype=torch.uint8).unsqueeze(1).unsqueeze(2))
        validx = (idx.sum(0) == 3) 
        masks[validx] = torch.tensor(mapping[k], dtype=torch.long)
    return masks
    


Below is my predict_mask function which i am not able to finish it up-
def predict_mask(img_path,model):
    img = Image.open(img_path)
    
    transform = transforms.Compose([transforms.Resize((512,512)),
                                    transforms.ToTensor()])
    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)
    return output

Right now, it is returning the tensor of shape torch.Size([1, 4, 512, 512]), how to convert it in the required format which should be a rgb image?