From dd8d6c70f074a931858f0940f9e0853fe4e1d79d Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Thu, 12 Dec 2024 09:28:44 -0800 Subject: [PATCH] Update [ghstack-poisoned] --- tensordict/utils.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tensordict/utils.py b/tensordict/utils.py index ae6d3b44f..6a3184d83 100644 --- 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 ):