fix wrapper order

pyro-ppl · Feb 8, 2024 · a64c092 · a64c092
1 parent ef3a3f7
commit a64c092
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/numpyro/optim.py b/numpyro/optim.py
@@ -36,10 +36,10 @@
 
 def _value_and_grad(f, x, forward_mode_differentiation=False):
     if forward_mode_differentiation:
-        def _wrapper(h, x):
-            out, aux = h(x)
+        def _wrapper(x):
+            out, aux = f(x)
             return out, (out, aux)
-        grads, (out, aux) = _wrapper(jacfwd(f, has_aux=True), x)
+        grads, (out, aux) = jacfwd(_wrapper, has_aux=True)(x)
         return (out, aux), grads
     else:
         return value_and_grad(f, has_aux=True)(x)