Update

[ghstack-poisoned]
pytorch · Dec 12, 2024 · dd8d6c7 · dd8d6c7
1 parent 97cd086
commit dd8d6c7
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tensordict/utils.py b/tensordict/utils.py
@@ -1548,6 +1548,11 @@ def assert_close(
 
     from tensordict._lazy import LazyStackedTensorDict
 
+    if is_tensorclass(actual):
+        actual = actual._tensordict
+    if is_tensorclass(expected):
+        expected = expected._tensordict
+
     if isinstance(actual, LazyStackedTensorDict) and isinstance(
         expected, LazyStackedTensorDict
     ):