Update

[ghstack-poisoned]
pytorch · Oct 25, 2024 · eb4d92a · eb4d92a
1 parent dd5eb67
commit eb4d92a
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/tensordict/utils.py b/tensordict/utils.py
@@ -1511,12 +1511,8 @@ def assert_close(
         return True
 
     try:
-        set1 = set(
-            actual.keys(is_leaf=lambda x: not is_non_tensor(x), leaves_only=True)
-        )
-        set2 = set(
-            expected.keys(is_leaf=lambda x: not is_non_tensor(x), leaves_only=True)
-        )
+        set1 = set(actual.keys())
+        set2 = set(expected.keys())
     except ValueError:
         # Persistent tensordicts do not work with is_leaf
         set1 = set(actual.keys(is_leaf=lambda cls: issubclass(cls, torch.Tensor)))