fix rebase functional-autograd-checkpoint

Lightning-AI · Oct 18, 2024 · 63e426d · 63e426d
1 parent 1a454f1
commit 63e426d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/thunder/torch/__init__.py b/thunder/torch/__init__.py
@@ -5262,8 +5262,8 @@ def _backward_checkpoint(
 ) -> tuple[None | TensorLike, ...]:
     from thunder.core.transforms import vjp
 
-    result = vjp(function)(args, grad_outputs, **kwargs)
-    return result
+    result, grads = vjp(function)(args, grad_outputs, **kwargs)
+    return grads
 
 
 #