Skip to content

Commit 8690ae7

Browse files
AlexanderSaninclaudeericspod
authored
fix(MetaTensor): astype with torch dtype now returns MetaTensor preserving metadata (#8911)
## Summary - `MetaTensor.astype()` called with a torch dtype (e.g. `torch.int32`, `torch.float16`) was silently returning a plain `torch.Tensor`, discarding all metadata (affine matrix, spacing, applied_operations, and any custom keys). - Root cause: `out_type` was hardcoded to `torch.Tensor` instead of `type(self)` (`MetaTensor`), so `convert_data_type` set `track_meta=False` and stripped the metadata. - Fix: use `out_type = type(self)` when `mod_str == "torch"`, so `convert_data_type` receives `output_type=MetaTensor`, sets `track_meta=True`, and the dtype cast is performed while preserving all metadata. - The `auto3dseg/analyzer.py` module already annotated `label_tensor.astype(torch.int16)` as returning a `MetaTensor` (line 493), relying on this contract. Closes #8202 ## Test plan - [ ] Existing `test_astype` test updated to assert `isinstance(result, MetaTensor)` and that metadata keys survive the cast. - [ ] All 96 `tests/data/meta_tensor/` tests pass locally (0 failures). - [ ] Manual verification: ```python import torch from monai.data import MetaTensor t = MetaTensor(torch.tensor([1., 2., 3.]), meta={"fname": "scan.nii"}) result = t.astype(torch.int32) assert isinstance(result, MetaTensor) # was torch.Tensor before assert result.meta["fname"] == "scan.nii" # metadata preserved assert result.dtype == torch.int32 # dtype correctly cast ``` Signed-off-by: Oleksandr Sanin <alexaaander.sanin@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com>
1 parent ed76cd5 commit 8690ae7

2 files changed

Lines changed: 9 additions & 4 deletions

File tree

monai/data/meta_tensor.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -495,7 +495,8 @@ def astype(self, dtype, device=None, *_args, **_kwargs):
495495
_kwargs: additional kwargs (currently unused).
496496
497497
Returns:
498-
data array instance
498+
``MetaTensor`` when a torch dtype is given (metadata is preserved),
499+
or ``np.ndarray`` when a numpy dtype is given.
499500
"""
500501
if isinstance(dtype, str):
501502
mod_str, *dtype = dtype.split(".", 1)
@@ -506,7 +507,7 @@ def astype(self, dtype, device=None, *_args, **_kwargs):
506507

507508
out_type: type[torch.Tensor] | type[np.ndarray] | None
508509
if mod_str == "torch":
509-
out_type = torch.Tensor
510+
out_type = type(self)
510511
elif mod_str in ("numpy", "np"):
511512
out_type = np.ndarray
512513
else:

tests/data/meta_tensor/test_meta_tensor.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -435,8 +435,12 @@ def test_astype(self):
435435
for np_types in ("float32", "np.float32", "numpy.float32", np.float32, float, "int", np.uint16):
436436
self.assertIsInstance(t.astype(np_types), np.ndarray)
437437
for pt_types in ("torch.float", torch.float, "torch.float64"):
438-
self.assertIsInstance(t.astype(pt_types), torch.Tensor)
439-
self.assertIsInstance(t.astype("torch.float", device="cpu"), torch.Tensor)
438+
result = t.astype(pt_types)
439+
self.assertIsInstance(result, MetaTensor)
440+
self.assertEqual(result.meta.get("fname"), "filename")
441+
result = t.astype("torch.float", device="cpu")
442+
self.assertIsInstance(result, MetaTensor)
443+
self.assertEqual(result.meta.get("fname"), "filename")
440444

441445
def test_transforms(self):
442446
key = "im"

0 commit comments

Comments
 (0)