Fix neural SDE example for 0.7.0 release

Author: patrick-kidger

examples/neural_sde.ipynb

@@ -342,7 +342,7 @@
     "        return mu * t - theta * y\n",
     "\n",
     "    def diffusion(t, y, args):\n",
-    "        return 2 * sigma * t / t1\n",
+    "        return jnp.array([2 * sigma * t / t1])\n",
     "\n",
     "    bm = diffrax.UnsafeBrownianPath(shape=(), key=bm_key)\n",
     "    drift = diffrax.ODETerm(drift)\n",