[BugFix] Fix compilation of TC with non-tensor + batch-size + device (#1337)

Author: vmoens

tensordict/tensorclass.py

@@ -1275,7 +1275,6 @@ def wrapper(
             new_params.append(
                 inspect.Parameter("names", inspect.Parameter.KEYWORD_ONLY, default=None)
             )
-
     wrapper.__signature__ = init_sig.replace(parameters=params + new_params)
 
     return wrapper
@@ -2284,7 +2283,9 @@ def set_tensor(
                 self._non_tensordict[key] = value
                 return self
             if non_tensor:
-                value = NonTensorData(value)
+                value = NonTensorData(
+                    value, batch_size=self.batch_size, device=self.device
+                )
             if key in self._non_tensordict:
                 del self._non_tensordict[key]
             # Avoiding key clash, honoring the user input to assign tensor type data to the key

test/test_compile.py

@@ -36,6 +36,8 @@
 
 from tensordict.nn.functional_modules import _exclude_td_from_pytree
 
+from tensordict.tensorclass import TensorClass
+
 from torch.utils._pytree import SUPPORTED_NODES, tree_map
 
 TORCH_VERSION = version.parse(version.parse(torch.__version__).base_version)
@@ -1207,6 +1209,56 @@ def test_td_input_non_tdmodule_nontensor(self, compiled):
                 func(torch.zeros(()), 2.0)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda is not available")
+class TestCompileNontensor:
+    # Same issue with the decorator @tensorclass version
+    @pytest.fixture(scope="class")
+    def data(self):
+        return torch.zeros((4, 3), device="cuda")
+
+    class TensorClassWithNonTensorData(TensorClass["nocast"]):
+        tensor: torch.Tensor
+        non_tensor_data: int
+
+    def fn_no_device_no_batch_size(self, data):
+        a = self.TensorClassWithNonTensorData(tensor=data, non_tensor_data=1)
+        return a.tensor
+
+    def fn_no_device(self, data):
+        a = self.TensorClassWithNonTensorData(
+            tensor=data, non_tensor_data=1, batch_size=[4]
+        )
+        return a.tensor
+
+    def fn_with_device(self, data):
+        a = self.TensorClassWithNonTensorData(
+            tensor=data, non_tensor_data=1, batch_size=[4], device="cuda"
+        )
+        return a.tensor
+
+    def fn_with_device_without_batch_size(self, data):
+        a = self.TensorClassWithNonTensorData(
+            tensor=data, non_tensor_data=1, device="cuda"
+        )
+        return a.tensor
+
+    def test_nontensor_no_device_no_batch_size(self, data):
+        torch._dynamo.reset_code_caches()
+        torch.compile(self.fn_no_device_no_batch_size)(data)
+
+    def test_nontensor_no_device(self, data):
+        torch._dynamo.reset_code_caches()
+        torch.compile(self.fn_no_device)(data)
+
+    def test_nontensor_with_device(self, data):
+        torch._dynamo.reset_code_caches()
+        torch.compile(self.fn_with_device)(data)
+
+    def test_nontensor_with_device_without_batch_size(self, data):
+        torch._dynamo.reset_code_caches()
+        torch.compile(self.fn_with_device_without_batch_size)(data)
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)