tidied up neural ODE example

Author: patrick-kidger

docs/examples/neural_ode.ipynb

@@ -198,7 +198,7 @@
     "def main(\n",
     "    dataset_size=256,\n",
     "    batch_size=32,\n",
-    "    lr_strategy=(3e-3, 3e-3),\n",
+    "    lr=3e-3,\n",
     "    steps_strategy=(500, 500),\n",
     "    length_strategy=(0.1, 1),\n",
     "    width_size=64,\n",
@@ -214,6 +214,7 @@
     "    _, length_size, data_size = ys.shape\n",
     "\n",
     "    model = NeuralODE(data_size, width_size, depth, key=model_key)\n",
+    "    optim = optax.adabelief(lr)\n",
     "\n",
     "    # Training loop like normal.\n",
     "    #\n",
@@ -233,8 +234,7 @@
     "        model = eqx.apply_updates(model, updates)\n",
     "        return loss, model, opt_state\n",
     "\n",
-    "    for lr, steps, length in zip(lr_strategy, steps_strategy, length_strategy):\n",
-    "        optim = optax.adabelief(lr)\n",
+    "    for steps, length in zip(steps_strategy, length_strategy):\n",
     "        opt_state = optim.init(eqx.filter(model, eqx.is_inexact_array))\n",
     "        _ts = ts[: int(length_size * length)]\n",
     "        _ys = ys[:, : int(length_size * length)]\n",
@@ -326,9 +326,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "jax0227",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
-   "name": "jax0227"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -340,7 +340,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.7"
+   "version": "3.11.8"
   }
  },
  "nbformat": 4,