amend

pytorch · Jan 18, 2024 · 25e87d3 · 25e87d3
1 parent be802e9
commit 25e87d3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tensordict/nn/utils.py b/tensordict/nn/utils.py
@@ -307,7 +307,7 @@ def __init__(self, mode):
         self.mode = mode
 
     def __call__(self, func):
-        @wraps(func)
+        @functools.wraps(func)
         def new_func(*args, **kwargs):
             with self:
                 return func(*args, **kwargs)