Fix introductory CNN example in recent JAX.

patrick-kidger · Oct 14, 2024 · d9b3ffd · d9b3ffd
1 parent 40eac53
commit d9b3ffd
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/examples/mnist.ipynb b/examples/mnist.ipynb
@@ -594,7 +594,9 @@
  " y: Int[Array, \" batch\"],\n",
  " ):\n",
  " loss_value, grads = eqx.filter_value_and_grad(loss)(model, x, y)\n",
- " updates, opt_state = optim.update(grads, opt_state, model)\n",
+ " updates, opt_state = optim.update(\n",
+ " grads, opt_state, eqx.filter(model, eqx.is_array)\n",
+ " )\n",
  " model = eqx.apply_updates(model, updates)\n",
  " return model, opt_state, loss_value\n",
  "\n",