Update thunder/executors/torchex.py

Lightning-AI · Oct 31, 2024 · 4252d4a · 4252d4a
1 parent f3a0857
commit 4252d4a
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/thunder/executors/torchex.py b/thunder/executors/torchex.py
@@ -171,6 +171,7 @@ def no_autocast(fn):
     prev_cpu = torch.is_autocast_cpu_enabled()
     prev = torch.is_autocast_enabled()
 
+    @wraps(fn)
     def no_autocast_fn(*args, **kwargs):
         try:
             set_autocast_enabled("cpu", False)