Fix Neural SDE example -- we changed how we check term compatibility.

patrick-kidger · Aug 8, 2024 · c8e1506 · c8e1506
1 parent d6d09dc
commit c8e1506
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/examples/neural_sde.ipynb b/examples/neural_sde.ipynb
@@ -135,6 +135,7 @@
     "        )\n",
     "\n",
     "    def __call__(self, t, y, args):\n",
+    "        t = jnp.asarray(t)\n",
     "        return self.scale * self.mlp(jnp.concatenate([t[None], y]))\n",
     "\n",
     "\n",
@@ -168,6 +169,7 @@
     "        self.hidden_size = hidden_size\n",
     "\n",
     "    def __call__(self, t, y, args):\n",
+    "        t = jnp.asarray(t)\n",
     "        return self.scale * self.mlp(jnp.concatenate([t[None], y])).reshape(\n",
     "            self.hidden_size, self.control_size\n",
     "        )"