diff --git a/src/torchjd/autogram/_jacobian_computer.py b/src/torchjd/autogram/_jacobian_computer.py
index 6bddbd60..a93bd5aa 100644
--- a/src/torchjd/autogram/_jacobian_computer.py
+++ b/src/torchjd/autogram/_jacobian_computer.py
@@ -195,5 +195,5 @@ def vmap(
         return cast(Matrix, jacobian), None
 
     @staticmethod
-    def setup_context(*_) -> None:
+    def setup_context(*_, **__) -> None:
         pass