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>
AlexanderSanin
requested review from
KumoLiu,
Nic-Ma and
ericspod
as code owners
|
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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
There was a problem hiding this comment.
Hi @AlexanderSanin I think this is a valid issue to consider here.
There was a problem hiding this comment.
@ericspod Agreed, thanks for confirming. Leaving this open for @AlexanderSanin to update the first_key == () branch to just return d (or restrict to key_iterator) so allow_missing_keys=True cases don't hit whole-dict conversion again.
|
Hi @AlexanderSanin I think this change is correct but please look at the coderabbit comment. Are there other transforms with this behaviour as well? It might be incorrect everywhere for the same reasons if so. Thanks! |
2 participants
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