Fixed 2nd device bug in teaser example

2025-03-31 16:25:31 +02:00 · 2025-03-31 16:25:31 +02:00 · 4919e7a429
commit 4919e7a429
parent cf13364482
1 changed files with 1 additions and 1 deletions
--- a/intro-teaser.ipynb
+++ b/intro-teaser.ipynb
@ -662,7 +662,7 @@
   "source": [
    "# Results\n",
    "plt.plot(X,Y,'.',label='Datapoints', color=\"lightgray\")\n",
-    "plt.plot(trajectory[-1][:,0].detach(),trajectory[-1][:,1].detach(),'.',label='Flow Matching', color=\"orange\") \n",
+    "plt.plot(trajectory[-1][:,0].detach().cpu(), trajectory[-1][:,1].detach().cpu(), '.',label='Flow Matching', color=\"orange\") \n",
    "plt.xlabel('x')\n",
    "plt.ylabel('y')\n",
    "plt.title('Probabilistic Version')\n",