From e943f9ab7657cd0dfca4fd1d91f9e96ff927c2ab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Val=C3=A9rian=20Rey?= Date: Wed, 14 Jan 2026 16:57:38 +0100 Subject: [PATCH] test: Fix using old backward interface in autograd_forward_backward --- tests/utils/forward_backwards.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/utils/forward_backwards.py b/tests/utils/forward_backwards.py index 7baabdb7..b0451e16 100644 --- a/tests/utils/forward_backwards.py +++ b/tests/utils/forward_backwards.py @@ -11,6 +11,7 @@ from torchjd.aggregation import Aggregator, Weighting from torchjd.autogram import Engine from torchjd.autojac import backward +from torchjd.autojac._jac_to_grad import jac_to_grad def autograd_forward_backward( @@ -29,7 +30,8 @@ def autojac_forward_backward( aggregator: Aggregator, ) -> None: losses = forward_pass(model, inputs, loss_fn, reduce_to_vector) - backward(losses, aggregator=aggregator) + backward(losses) + jac_to_grad(model.parameters(), aggregator) def autograd_gramian_forward_backward(