diff --git a/numpyro/optim.py b/numpyro/optim.py
index 37d16dc4a..5ad451bab 100644
--- 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)