len

Signed-off-by: Matteo Bettini <matbet@meta.com>
pytorch · vmoens · Aug 4, 2023 · Jul 6, 2023 · Jul 6, 2023 · Jul 6, 2023
commit 64cbe1b3138464f4a2c138af0a54c27d8cfb8bbf
diff --git a/test/test_specs.py b/test/test_specs.py
@@ -2530,6 +2530,10 @@ def test_unsqueeze(self):
         cus = cu.squeeze(0)
         assert cus == c
 
+    def test_len(self):
+        c = self._get_het_specs()
+        assert len(c) == 2
+
 
 # MultiDiscreteTensorSpec: Pending resolution of https://github.com/pytorch/pytorch/issues/100080.
 @pytest.mark.parametrize(

diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -3156,6 +3156,7 @@ def to_numpy(self, val: TensorDict, safe: bool = None) -> dict:
         return {key: self[key].to_numpy(val) for key, val in val.items()}
 
     def __len__(self):
+        """Returns the number of keys present in all the members of the LazyStack."""
         return len(self.keys())
 
     def values(