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: 96bfda0

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.