Skip to content

fix(RandGridDistortiond): only convert transform keys when skipping transform #8920

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
AlexanderSanin wants to merge 1 commit into
base: dev
Choose a base branch
from
+13 −3
Open

fix(RandGridDistortiond): only convert transform keys when skipping transform #8920

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions monai/transforms/spatial/dictionary.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2305,12 +2305,13 @@ def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> dict[Hashable, torc
d = dict(data)
self.randomize(None)
if not self._do_transform:
out: dict[Hashable, torch.Tensor] = convert_to_tensor(d, track_meta=get_track_meta())
return out
for key in self.key_iterator(d):
d[key] = convert_to_tensor(d[key], track_meta=get_track_meta())
return d

first_key: Hashable = self.first_key(d)
if first_key == ():
out = convert_to_tensor(d, track_meta=get_track_meta())
out: dict[Hashable, torch.Tensor] = convert_to_tensor(d, track_meta=get_track_meta())
return out
Comment on lines 2313 to 2315

@coderabbitai coderabbitai Bot Jun 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Avoid whole-dict tensor coercion when no configured key exists.

Line 2314 still converts the entire mapping when first_key == (). With allow_missing_keys=True, this can reintroduce non-key coercion and the same downstream collation failure this PR fixes. Return d unchanged (or only iterate key_iterator, which is empty here).

Suggested patch 
         first_key: Hashable = self.first_key(d)
         if first_key == ():
-            out: dict[Hashable, torch.Tensor] = convert_to_tensor(d, track_meta=get_track_meta())
-            return out
+            return d

As per coding guidelines, "Examine code for logical error or inconsistencies, and suggest what may be changed to addressed these."

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@monai/transforms/spatial/dictionary.py` around lines 2313 - 2315, When
first_key equals an empty tuple, the code incorrectly converts the entire
dictionary to tensors using convert_to_tensor, which causes unwanted non-key
coercion when allow_missing_keys is True and leads to collation failures. In the
if first_key == () block, simply return d unchanged instead of calling
convert_to_tensor(d, track_meta=get_track_meta()), since an empty first_key
means there are no configured keys to process.

Source: Coding guidelines

if isinstance(d[first_key], MetaTensor) and d[first_key].pending_operations: # type: ignore
warnings.warn(f"data['{first_key}'] has pending operations, transform may return incorrect results.")
Expand Down
9 changes: 9 additions & 0 deletions tests/transforms/test_rand_grid_distortiond.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,15 @@ def test_rand_grid_distortiond(self, input_param, seed, input_data, expected_val
assert_allclose(result["img"], expected_val_img, type_test=False, rtol=1e-4, atol=1e-4)
assert_allclose(result["mask"], expected_val_mask, type_test=False, rtol=1e-4, atol=1e-4)

def test_no_transform_preserves_non_image_keys(self):
"""Non-image dict entries must not be coerced to tensors when the transform is skipped."""
img = np.indices([6, 6]).astype(np.float32)
data = {"img": img, "label": 42, "filename": "scan.nii"}
g = RandGridDistortiond(keys=["img"], prob=0.0)
result = g(data)
self.assertIsInstance(result["label"], int)
self.assertIsInstance(result["filename"], str)


if __name__ == "__main__":
unittest.main()
Loading