[Test] Test non-tensor propagation in arithmetic ops (#926)

pytorch · Jul 30, 2024 · 484a045 · 484a045 · github-actions · Jul 30, 2024
1 parent 378c53f
commit 484a045
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
@@ -2589,6 +2589,11 @@ def test_add_default(self):
         assert "d" not in tdadd
         assert "b" in tdadd
 
+        td0 = TensorDict(a=1, b=1, c=1, non_tensor="a string")
+        td1 = TensorDict(b=2, c=2, d=2, non_tensor="a string")
+        td = td0.add(td1, default=torch.zeros(()))
+        assert td["non_tensor"] == "a string"
+
     def test_sub_default(self):
         # Create two tds with different key sets
         td0 = TensorDict(a=1, b=1, c=1)