title: Pytorch: Custom dataset for segmentation, error with input dimension shape
body: <p>I would like to know why I get this error: <code>1only batches of spatial targets supported (3D tensors) but got targets of size  [16, 3, 512, 512]</code>, with this dataset:</p>
  <pre><code>class MyDataset (Dataset):
  def __init__(self, images, masks, mean, std, transforms=None):
      self.images = images
      self.masks = masks
      self.transforms = transforms
      self.mean = mean
      self.std = std
  
  def __len__(self):
      return len(self.images)
  
  def __getitem__(self, idx):
      image_path = self.images[idx]
      image = Image.open(image_path)
      image = np.array(image)
      mask_path = self.masks[idx]
      mask = Image.open(mask_path)
      mask = np.array(mask)
      
      
      if self.transforms is not None:
          aug = self.transforms(image=image, mask=mask)
          image = T.ToPILImage()(aug['image'])
          mask = aug['mask']
          
      
      if self.transforms is None:
          image = T.ToPILImage()(image)
      
      t = T.Compose([T.Resize(256), T.ToTensor(), T.Normalize(self.mean, self.std)])
      image = t(image)
      mask = torch.from_numpy(mask).long().permute(2, 0, 1)
      
      return image, mask
  </code></pre>
  <p>I have checked that images and masks have shapes [C, H, W], and I don't know why the mini-batch dimension is added as first parameter.</p>