[BugFix] Better comparison of tensorclasses

ghstack-source-id: 8def6f01f2b6d09714319a56f96b166ac1fd49d5 Pull Request resolved: #1137
pytorch · Dec 12, 2024 · eb4a56e · eb4a56e · github-actions · Dec 12, 2024
1 parent d7506c4
commit eb4a56e
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
     ):