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?