Code update

PiperOrigin-RevId: 610472835
google-research · Feb 27, 2024 · 7c13a89 · 7c13a89
1 parent 585a2af
commit 7c13a89
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/swirl_dynamics/projects/probabilistic_diffusion/colabs/demo.ipynb b/swirl_dynamics/projects/probabilistic_diffusion/colabs/demo.ipynb
@@ -973,7 +973,7 @@
         "num_samples_per_cond = 5\n",
         "\n",
         "generate = jax.jit(\n",
-        "    functools.partial(sampler.generate, num_samples_per_cond)\n",
+        "    functools.partial(cond_sampler.generate, num_samples_per_cond)\n",
         ")"
       ]
     },