[BugFix] NonTensor.encode must return a NonTensorData

ghstack-source-id: 36fc87c
Pull-Request-resolved: #2944

Author: Vincent Moens

test/test_specs.py

@@ -3940,10 +3940,17 @@ def test_device_cast(self):
         assert comp["nontensor"].device == torch.device("cpu")
 
     def test_encode(self):
+        nt = NonTensor(shape=(1,))
+        r = nt.encode("a string")
+        assert isinstance(r, NonTensorData)
+        assert r.shape == nt.shape
+
         comp = Composite(device="cpu")
-        comp["nontensor"] = NonTensor(shape=())
+        comp["nontensor"] = nt
         r = comp.encode({"nontensor": "a string"})
-        assert isinstance(r["nontensor"], str)
+        assert isinstance(r, TensorDict)
+        assert isinstance(r.get("nontensor"), NonTensorData)
+        assert r.get("nontensor").shape == (1,)
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="not cuda device")

torchrl/data/tensor_specs.py

@@ -2918,7 +2918,7 @@ def _encode_eager(
         *,
         ignore_device: bool = False,
     ) -> torch.Tensor | TensorDictBase:
-        return val
+        return NonTensorData(val, device=self.device, batch_size=self.shape)
 
 
 class _UnboundedMeta(abc.ABCMeta):