fix(RandGridDistortiond): only convert transform keys when skipping transform#8920
Conversation
…ransform When `_do_transform` is False, `RandGridDistortiond.__call__` was calling `convert_to_tensor(d, ...)` on the entire input dict. This recursively converts *all* values — including non-image entries such as integers and scalars — into PyTorch tensors. The converted dict is then returned to the DataLoader which hands it to MONAI's `collate_meta_tensor_fn`. That collate path expects non-image entries to remain as their original Python types; receiving 0-d tensors instead triggers an `AttributeError: 'int' object has no attribute 'numel'` when the collate function iterates over what it believes to be a batch of tensors. Fix: iterate over `self.key_iterator(d)` and convert only those values, exactly as the transform loop further down in the same method already does. This matches the per-key pattern used in sibling transforms such as `RandAffined` and leaves unrelated dict entries unchanged. Also adds a regression test that verifies integer and string entries are preserved when the transform is skipped (prob=0.0). Closes Project-MONAI#8604 Signed-off-by: Oleksandr Sanin <alexaaander.sanin@gmail.com>
|
Hey @ericspod @garciadias. Could you, please, have a look at this? |
📝 WalkthroughWalkthrough
Estimated code review effort🎯 2 (Simple) | ⏱️ ~5 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/transforms/test_rand_grid_distortiond.py (1)
88-96: ⚡ Quick winStrengthen regression test with keyed-value conversion assertion.
The new test verifies non-key preservation, but it should also assert the keyed
"img"is still tensor-converted in the no-op path to lock the full contract.Suggested patch
import numpy as np +import torch from parameterized import parameterized @@ result = g(data) + self.assertIsInstance(result["img"], torch.Tensor) self.assertIsInstance(result["label"], int) self.assertIsInstance(result["filename"], str)As per coding guidelines, "Ensure new or modified definitions will be covered by existing or new unit tests."
🤖 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 `@tests/transforms/test_rand_grid_distortiond.py` around lines 88 - 96, The test_no_transform_preserves_non_image_keys method verifies that non-keyed entries are preserved but does not validate that the keyed entry "img" is still properly converted to a tensor when the RandGridDistortiond transform is skipped due to prob=0.0. Add assertions after the g(data) call to verify that result["img"] is properly converted to a tensor type, ensuring the full contract of the transform is tested including the tensor conversion behavior for keyed values even in the no-op probability path.Source: Coding guidelines
🤖 Prompt for all review comments with 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.
Inline comments:
In `@monai/transforms/spatial/dictionary.py`:
- Around line 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.
---
Nitpick comments:
In `@tests/transforms/test_rand_grid_distortiond.py`:
- Around line 88-96: The test_no_transform_preserves_non_image_keys method
verifies that non-keyed entries are preserved but does not validate that the
keyed entry "img" is still properly converted to a tensor when the
RandGridDistortiond transform is skipped due to prob=0.0. Add assertions after
the g(data) call to verify that result["img"] is properly converted to a tensor
type, ensuring the full contract of the transform is tested including the tensor
conversion behavior for keyed values even in the no-op probability path.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: a6e1a8bf-052e-4610-a1c7-241311fcb4c5
📒 Files selected for processing (2)
monai/transforms/spatial/dictionary.pytests/transforms/test_rand_grid_distortiond.py
| 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 |
There was a problem hiding this comment.
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 dAs 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
Description
Fixes #8604
Root Cause
RandGridDistortiond.__call__usesconvert_to_tensor(d, track_meta=...)on the entire input dict when the transform is skipped (_do_transform=False).convert_to_tensorrecursively converts every value in the dict — including non-image entries such as scalars and integers — into PyTorch tensors.When the DataLoader collates a batch of these dicts, MONAI's
collate_meta_tensor_fnis invoked for entries that now appear as tensors. That collate path expects non-image keys to retain their original Python types; receiving 0-d tensors instead triggers:Fix
Replace the whole-dict conversion with a per-key loop using
self.key_iterator(d), so only the keys this transform is responsible for are converted. This matches the pattern already used in thefor key, mode, padding_mode in self.key_iterator(...)loop further down in the same method, and mirrors the approach in sibling transforms such asRandAffined.Changes
monai/transforms/spatial/dictionary.py: usekey_iteratorin the no-op branch ofRandGridDistortiond.__call__tests/transforms/test_rand_grid_distortiond.py: add regression test asserting that non-image dict entries (integer, string) are preserved when the transform is skippedTest plan
test_no_transform_preserves_non_image_keysconfirms integer and string dict entries survive a skipped transformtest_rand_grid_distortiond_0/1/2) still passDataLoaderwithRandGridDistortiond(prob=0.0)no longer raisesAttributeErrorwhen the batch contains integer metadata fields alongside image tensors