Update score_based_diffusion.ipynb

patrick-kidger · Jun 8, 2024 · 262f347 · 262f347
1 parent 1f07888
commit 262f347
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/score_based_diffusion.ipynb b/examples/score_based_diffusion.ipynb
@@ -158,7 +158,7 @@
     "        self.t1 = t1\n",
     "\n",
     "    def __call__(self, t, y):\n",
-    "        t = t / self.t1\n",
+    "        t = jnp.array(t / self.t1)\n",
     "        _, height, width = y.shape\n",
     "        t = einops.repeat(t, \"-> 1 h w\", h=height, w=width)\n",
     "        y = jnp.concatenate([y, t])\n",