Update

[ghstack-poisoned]
pytorch · Oct 18, 2024 · 1787e59 · 1787e59
1 parent ccb4dcc
commit 1787e59
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/tensordict/tensorclass.py b/tensordict/tensorclass.py
@@ -1145,11 +1145,8 @@ def wrapper(self, key: str, value: Any) -> None:  # noqa: D417
 
 def _wrap_td_method(funcname, *, copy_non_tensor=False, no_wrap=False):
     def check_out(kwargs, result):
-        out = kwargs.get("out")
-        if out is result:
-            # No need to transform output
-            return True
-        return False
+        # No need to transform output if True
+        return kwargs.get("out") is result
 
     def deliver_result(self, result, kwargs):
         if result is None: