AttributeError in DataLoader when using RandGridDistortiond transform

[albertojbc]

Description I have tried to include RandGridDistortiond transform as part of my data augmentation pipeline. However, when I include this transform the following error raises:

Exception has occurred: AttributeError
Caught AttributeError in DataLoader worker process 0.
AttributeError: Caught AttributeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "path/lib/python3.10/site-packages/torch/utils/data/_utils/worker.py", line 351, in _worker_loop
    data = fetcher.fetch(index)  # type: ignore[possibly-undefined]
  File "path/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 55, in fetch
    return self.collate_fn(data)
  File "path/lib/python3.10/site-packages/monai/data/utils.py", line 514, in list_data_collate
    ret[key] = collate_fn(data_for_batch)
  File "path/lib/python3.10/site-packages/torch/utils/data/_utils/collate.py", line 398, in default_collate
    return collate(batch, collate_fn_map=default_collate_fn_map)
  File "path/lib/python3.10/site-packages/torch/utils/data/_utils/collate.py", line 155, in collate
    return collate_fn_map[elem_type](batch, collate_fn_map=collate_fn_map)
  File "path/lib/python3.10/site-packages/monai/data/utils.py", line 458, in collate_meta_tensor_fn
    collated = collate_fn(batch)  # type: ignore
  File "path/lib/python3.10/site-packages/torch/utils/data/_utils/collate.py", line 269, in collate_tensor_fn
    numel = sum(x.numel() for x in batch)
  File "path/lib/python3.10/site-packages/torch/utils/data/_utils/collate.py", line 269, in <genexpr>
    numel = sum(x.numel() for x in batch)
AttributeError: 'int' object has no attribute 'numel'

Reproduction details This is how I defined the transform:

RandGridDistortiond(keys=["image"], prob=0.3, distort_limit=(-0.35, 0.35), padding_mode='zeros'),

Environment details Python 3.10.16 MONAI version: 1.3.1 Numpy version: 1.24.3 Pytorch version: 2.5.0+cu124 MONAI flags: HAS_EXT = False, USE_COMPILED = False, USE_META_DICT = False MONAI rev id: 96bfda00c6bd290297f5e3514ea227c6be4d08b4

Optional dependencies: Pytorch Ignite version: 0.4.11 ITK version: 5.3.0 Nibabel version: 5.1.0 scikit-image version: 0.22.0 scipy version: 1.11.1 Pillow version: 9.5.0 Tensorboard version: 2.14.1 gdown version: 4.7.1 TorchVision version: 0.20.0+cu124 tqdm version: 4.65.0 lmdb version: 1.4.1 psutil version: 5.9.0 pandas version: 2.0.3 einops version: 0.7.0 transformers version: 4.49.0 mlflow version: 2.7.1 pynrrd version: 1.0.0 clearml version: 1.13.1

Found solution I managed to solve this issue in my case by changing the __call__ function of RandGridDistortiond in monai/transforms/spatial/dictionary.py. Instead of:

if not self._do_transform:
      out: dict[Hashable, torch.Tensor] = convert_to_tensor(d, track_meta=get_track_meta())
      return out

Use this:

if not self._do_transform:
      for key in self.key_iterator(d):
          d[key] = convert_to_tensor(d[key], track_meta=get_track_meta())
      return d

For me, this change made the DataLoader work properly.