fix: optimistix deprecated API

wcxve · Oct 9, 2024 · 043638d · 043638d
1 parent 3450003
commit 043638d
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/elisa/infer/fit.py b/src/elisa/infer/fit.py
@@ -171,7 +171,8 @@ def lm(init):
                     max_steps=max_steps,
                     throw=throw,
                 )
-                return res.value, jnp.linalg.norm(res.state.f_info.grad)
+                grad_norm = jnp.linalg.norm(res.state.f_info.compute_grad())
+                return res.value, grad_norm
 
             self._lm = jax.jit(lm)