[Minor] Fix type deletion in tensorclass load_memmap (#602)

pytorch · Dec 21, 2023 · 5b6573a · 5b6573a
1 parent 306670a
commit 5b6573a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tensordict/tensorclass.py b/tensordict/tensorclass.py
@@ -412,8 +412,8 @@ def save_metadata(cls=cls, _non_tensordict=_non_tensordict, prefix=prefix):
 
 
 def _load_memmap(cls, prefix: Path, metadata: dict):
-    del metadata["_type"]
     non_tensordict = copy(metadata)
+    del non_tensordict["_type"]
     if os.path.exists(prefix / "other.pickle"):
         with open(prefix / "other.pickle", "rb") as pickle_file:
             non_tensordict.update(pickle.load(pickle_file))