fix vmapped_log_probs example

jax-ml · Jul 11, 2020 · 8c8869c · 8c8869c
1 parent 5bbc195
commit 8c8869c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/docs/notebooks/vmapped_log_probs.ipynb b/docs/notebooks/vmapped_log_probs.ipynb
@@ -415,7 +415,7 @@
         "    beta_sample = beta_loc + np.exp(beta_log_scale) * epsilon\n",
         "    return np.mean(batched_log_joint(beta_sample), 0) + np.sum(beta_log_scale - 0.5 * onp.log(2*onp.pi))\n",
         " \n",
-        "elbo = jax.jit(elbo, static_argnums=(1, 2))\n",
+        "elbo = jax.jit(elbo)\n",
         "elbo_val_and_grad = jax.jit(jax.value_and_grad(elbo, argnums=(0, 1)))"
       ],
       "execution_count": 0,
@@ -653,4 +653,4 @@
       "outputs": []
     }
   ]
-}
+}