From 1a10ba5e9e07700177191e0e1d2ecd80ffb903aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Val=C3=A9rian=20Rey?= Date: Tue, 9 Sep 2025 15:50:53 +0200 Subject: [PATCH] Add vmap rule generation to JacobianAccumulator --- src/torchjd/autogram/_module_hook_manager.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/torchjd/autogram/_module_hook_manager.py b/src/torchjd/autogram/_module_hook_manager.py index e5087bb0..07381184 100644 --- a/src/torchjd/autogram/_module_hook_manager.py +++ b/src/torchjd/autogram/_module_hook_manager.py @@ -108,6 +108,8 @@ class JacobianAccumulator(torch.autograd.Function): toggle mechanism to activate only during the Gramian accumulation phase. """ + generate_vmap_rule = True + @staticmethod def forward(*xs: Tensor) -> tuple[Tensor, ...]: return tuple([x.detach() for x in xs])