Hi everyone!

As i mentioned in the topic, I got an error with the permute function.

```
def __getitem__(self, i: int) -> Tuple[torch.Tensor, torch.Tensor]:
""" Input-Output Generator by index
Args:
i (int): Order of data
Returns:
(Tuple[torch.Tensor, torch.Tensor]): Image and mask PyTorch tensor
"""
_id = self.__image_ids[i]
paths = self.__id2path(_id)
image = self.__load_n_preprocess(paths["image_path"], False)
mask = self.__load_n_preprocess(paths["mask_path"], True)
image = torch.from_numpy(image)
mask = torch.from_numpy(mask)
image = image.permute(2, 0, 1)
mask = mask.permute(2, 0, 1)
return image, mask
```

ERROR

```
mask = mask.permute(2, 0, 1)
RuntimeError: number of dims don't match in permute
```

I realized that It gets stuck on mask.permute after doing image.permute just fine. By the way image’s shape is torch.Size([256, 256, 3]) mask’s shape is torch.Size([256, 256]).

Also, this error arised when this function was called the third time.