Update for loading MetaTensor

Signed-off-by: Eric Kerfoot <[email protected]>
Project-MONAI · Jan 29, 2025 · d6152a3 · d6152a3
1 parent a375354
commit d6152a3
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/monai/data/meta_tensor.py b/monai/data/meta_tensor.py
@@ -607,3 +607,6 @@ def print_verbose(self) -> None:
         print(self)
         if self.meta is not None:
             print(self.meta.__repr__())
+
+
+torch.serialization.add_safe_globals([MetaTensor])
diff --git a/tests/test_meta_tensor.py b/tests/test_meta_tensor.py
@@ -256,7 +256,7 @@ def test_pickling(self):
         with tempfile.TemporaryDirectory() as tmp_dir:
             fname = os.path.join(tmp_dir, "im.pt")
             torch.save(m, fname)
-            m2 = torch.load(fname)
+            m2 = torch.load(fname, weights_only=False)
             if not isinstance(m2, MetaTensor) and not pytorch_after(1, 8, 1):
                 warnings.warn("Old version of pytorch. pickling converts `MetaTensor` to `torch.Tensor`.")
                 m = m.as_tensor()