[BugFix] Fix setting non-tensors as data in NonTensorData (#864)

tensordict/tensorclass.py

@@ -2417,11 +2417,11 @@ def _stack_non_tensor(cls, list_of_non_tensor, dim=0):
         first = list_of_non_tensor[0]
 
         def _check_equal(a, b):
-            if isinstance(a, _ACCEPTED_CLASSES) or isinstance(b, _ACCEPTED_CLASSES):
-                return (a == b).all()
-            if isinstance(a, np.ndarray) or isinstance(b, np.ndarray):
-                return (a == b).all()
             try:
+                if isinstance(a, _ACCEPTED_CLASSES) or isinstance(b, _ACCEPTED_CLASSES):
+                    return (a == b).all() and a.shape == b.shape
+                if isinstance(a, np.ndarray) or isinstance(b, np.ndarray):
+                    return (a == b).all() and a.shape == b.shape
                 iseq = a == b
             except Exception:
                 iseq = False

test/test_tensordict.py

@@ -9127,6 +9127,16 @@ def test_nontensor_dict(self, non_tensor_data):
             TensorDict.from_dict(non_tensor_data.to_dict()) == non_tensor_data
         ).all()
 
+    def test_nontensor_tensor(self):
+        t1 = torch.tensor([1, 2, 3], dtype=torch.float)
+        t2 = torch.tensor([1, 2, 3, 4], dtype=torch.float)
+        stack = NonTensorStack(NonTensorData(t1), NonTensorData(t2))  # this works fine
+        assert all(isinstance(t, torch.Tensor) for t in stack.tolist())
+        stack = torch.stack(
+            [NonTensorData(t1), NonTensorData(t2)]
+        )  # this triggers an exception
+        assert all(isinstance(t, torch.Tensor) for t in stack.tolist())
+
     def test_set(self, non_tensor_data):
         non_tensor_data.set(("nested", "another_string"), "another string!")
         assert (