[BugFix] Track sub-tds in memmap (#719)

tensordict/_td.py

@@ -1884,7 +1884,8 @@ def save_metadata(data: TensorDictBase, filepath, metadata=None):
             dest._is_shared = False  # since they are mutually exclusive
 
         for key, value in self.items():
-            if _is_tensor_collection(value.__class__):
+            type_value = type(value)
+            if _is_tensor_collection(type_value):
                 dest._tensordict[key] = value._memmap_(
                     prefix=prefix / key if prefix is not None else None,
                     copy_existing=copy_existing,
@@ -1894,6 +1895,10 @@ def save_metadata(data: TensorDictBase, filepath, metadata=None):
                     like=like,
                     share_non_tensor=share_non_tensor,
                 )
+                if prefix is not None:
+                    metadata[key] = {
+                        "type": type_value.__name__,
+                    }
                 continue
             else:
                 # user did specify location and memmap is in wrong place, so we copy
@@ -1944,10 +1949,15 @@ def _load_memmap(cls, prefix: str, metadata: dict) -> T:
 
         out = cls({}, batch_size=metadata.pop("shape"), device=metadata.pop("device"))
 
+        paths = set()
         for key, entry_metadata in metadata.items():
             if not isinstance(entry_metadata, dict):
                 # there can be other metadata
                 continue
+            type_value = entry_metadata.get("type", None)
+            if type_value is not None:
+                paths.add(key)
+                continue
             dtype = entry_metadata.get("dtype", None)
             shape = entry_metadata.get("shape", None)
             if (
@@ -1970,7 +1980,7 @@ def _load_memmap(cls, prefix: str, metadata: dict) -> T:
             )
         # iterate over folders and load them
         for path in prefix.iterdir():
-            if path.is_dir():
+            if path.is_dir() and path.parts[-1] in paths:
                 key = path.parts[len(prefix.parts) :]
                 out.set(key, TensorDict.load_memmap(path))
         return out

test/test_tensordict.py

@@ -3101,6 +3101,8 @@ def test_memmap_(self, td_name, device, use_dir, tmpdir, num_threads):
             )
             assert td.is_memmap(), (td, td._is_memmap)
         if use_dir:
+            # This would fail if we were not filtering out unregistered sub-folders
+            os.mkdir(Path(tmpdir) / "some_other_path")
             assert_allclose_td(TensorDict.load_memmap(tmpdir), td)
 
     @pytest.mark.parametrize("copy_existing", [False, True])