RuntimeError: Caught RuntimeError in DataLoader worker process 0. &RuntimeError: Trying to resize storage that is not resizable

DataModule
1

I wrote a LightningDataModule for the MSCOCO dataset and implemented init, setup, train_dataloader, val_dataloader, and test_dataloader. And in the setup func, I instance a torch.utils.data.Dataset which will resize mscoco img to (128, 128) in the getitem func:

    def setup(self, stage: Optional[str] = None) -> None:
        """
        Creates train, val, and test dataset
        """
        if stage == "fit" or stage is None:
            self.dataset_train = self.dataset_cls(self.data_root, split=self.split, transform=self.transform, catnames=self.catnames)
            self.dataset_val = self.dataset_cls(self.data_root, split=self.split, transform=self.transform, catnames=self.catnames)

        if stage == "test" or stage is None:
            self.dataset_test = self.dataset_cls(self.data_root, split='val', transform=self.transform, catnames=self.catnames)

The self.transform=transforms.Compose([transforms.ToTensor(), transforms.Resize((128, 128))]), and the train_dataloader goes to:

    def train_dataloader(self, *args: Any, **kwargs: Any) -> DataLoader:
        """ The train dataloader """
        return self._data_loader(self.dataset_train, shuffle=self.shuffle)

    def _data_loader(self, dataset: Dataset, shuffle: bool = False) -> DataLoader:
        return DataLoader(
            dataset,
            batch_size=self.batch_size,
            shuffle=shuffle,
            num_workers=self.num_workers,
            drop_last=self.drop_last,
            pin_memory=self.pin_memory
        )

My torch.utils.data.Dataset works well with dataloader and can get the right output. But it has errors in the LightningDataModule:

RuntimeError: Caught RuntimeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/xxx/lib/python3.8/site-packages/torch/utils/data/_utils/worker.py", line 287, in _worker_loop
    data = fetcher.fetch(index)
  File "/xxx/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 52, in fetch
    return self.collate_fn(data)
  File "/xxx/lib/python3.8/site-packages/torch/utils/data/_utils/collate.py", line 172, in default_collate
    return [default_collate(samples) for samples in transposed]  # Backwards compatibility.
  File "xxx/lib/python3.8/site-packages/torch/utils/data/_utils/collate.py", line 172, in <listcomp>
    return [default_collate(samples) for samples in transposed]  # Backwards compatibility.
  File "xxx/lib/python3.8/site-packages/torch/utils/data/_utils/collate.py", line 137, in default_collate
    out = elem.new(storage).resize_(len(batch), *list(elem.size()))
RuntimeError: Trying to resize storage that is not resizable
2

I found the solution