[BugFix] Fix encode overrides

ghstack-source-id: 2b7710a
Pull-Request-resolved: #2943

Author: Vincent Moens

test/test_specs.py

@@ -3939,6 +3939,12 @@ def test_device_cast(self):
         comp["nontensor"] = NonTensor(device="cpu")
         assert comp["nontensor"].device == torch.device("cpu")
 
+    def test_encode(self):
+        comp = Composite(device="cpu")
+        comp["nontensor"] = NonTensor(shape=())
+        r = comp.encode({"nontensor": "a string"})
+        assert isinstance(r["nontensor"], str)
+
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="not cuda device")
 def test_device_ordinal():

torchrl/data/tensor_specs.py

@@ -1609,7 +1609,7 @@ def space(self):
     def _project(self, val: TensorDictBase) -> TensorDictBase:
         raise NOT_IMPLEMENTED_ERROR
 
-    def encode(
+    def _encode_eager(
         self, val: np.ndarray | torch.Tensor, *, ignore_device=False
     ) -> torch.Tensor:
         if self.dim != 0 and not isinstance(val, tuple):
@@ -2912,7 +2912,7 @@ def to_numpy(
     ) -> np.ndarray | dict:
         return val
 
-    def encode(
+    def _encode_eager(
         self,
         val: np.ndarray | torch.Tensor | TensorDictBase,
         *,
@@ -6357,7 +6357,7 @@ def empty(self):
             [spec.empty() for spec in self._specs], dim=self.stack_dim
         )
 
-    def encode(
+    def _encode_eager(
         self, vals: dict[str, Any], ignore_device: bool = False
     ) -> dict[str, torch.Tensor]:
         raise NOT_IMPLEMENTED_ERROR