fix bug

InternLM · Jul 9, 2024 · 4b41bb5 · 4b41bb5
1 parent 8951edb
commit 4b41bb5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/internlm/solver/activation_checkpoint.py b/internlm/solver/activation_checkpoint.py
@@ -62,7 +62,6 @@ def recompute_forward_context(args, no_communication):
 
         if handle:
             handle.wait()
-            args = list(args)
             args[0] = grad_output
 
 
@@ -157,6 +156,7 @@ def backward(ctx, *args):
 
         detached_inputs = detach_variable(tuple(inputs))
 
+        args = list(args)
         with recompute_forward_context(args, no_communication):
             if ctx.had_autocast_in_fwd:
                 with torch.enable_grad(), internlm_accelerator.amp.autocast():