diff --git a/docs/notebooks/2d-model_discriminant.ipynb b/docs/notebooks/2d-model_discriminant.ipynb
index be0e9c0b..de44c209 100644
--- a/docs/notebooks/2d-model_discriminant.ipynb
+++ b/docs/notebooks/2d-model_discriminant.ipynb
@@ -411,8 +411,8 @@
    "source": [
     "fig, ax = plt.subplots(figsize=(5,4),dpi=100)\n",
     "\n",
-    "loss_train = [x for x in model.loss_train]\n",
-    "loss_valid = [x for x in model.loss_valid]\n",
+    "loss_train = [x for x in model.logs['loss_train']]\n",
+    "loss_valid = [x for x in model.logs['loss_valid']]\n",
     "\n",
     "# Loss function\n",
     "ax.plot(loss_train,'-',label='Train')\n",
diff --git a/docs/notebooks/2d-model_tica.ipynb b/docs/notebooks/2d-model_tica.ipynb
index 21752996..640b6699 100644
--- a/docs/notebooks/2d-model_tica.ipynb
+++ b/docs/notebooks/2d-model_tica.ipynb
@@ -494,8 +494,8 @@
    "source": [
     "fig, ax = plt.subplots(figsize=(5,4),dpi=100)\n",
     "\n",
-    "loss_train = [x.cpu() for x in model.loss_train]\n",
-    "loss_valid = [x.cpu() for x in model.loss_valid]\n",
+    "loss_train = [x for x in model.logs['loss_train']]\n",
+    "loss_valid = [x for x in model.logs['loss_valid']]\n",
     "\n",
     "# Loss function\n",
     "ax.plot(loss_train,'-',label='Train')\n",
diff --git a/docs/notebooks/ala2_deeplda.ipynb b/docs/notebooks/ala2_deeplda.ipynb
index 624e0477..453a27a0 100644
--- a/docs/notebooks/ala2_deeplda.ipynb
+++ b/docs/notebooks/ala2_deeplda.ipynb
@@ -312,8 +312,8 @@
    "source": [
     "fig, ax = plt.subplots(figsize=(5,4),dpi=100)\n",
     "\n",
-    "loss_train = [x.cpu() for x in model.loss_train]\n",
-    "loss_valid = [x.cpu() for x in model.loss_valid]\n",
+    "loss_train = [x for x in model.logs['loss_train']]\n",
+    "loss_valid = [x for x in model.logs['loss_valid']]\n",
     "\n",
     "# Loss function\n",
     "ax.plot(loss_train,'-',label='Train')\n",
diff --git a/docs/notebooks/ala2_deeptica_multithermal.ipynb b/docs/notebooks/ala2_deeptica_multithermal.ipynb
index cc7234d3..3f064943 100644
--- a/docs/notebooks/ala2_deeptica_multithermal.ipynb
+++ b/docs/notebooks/ala2_deeptica_multithermal.ipynb
@@ -502,8 +502,8 @@
    "source": [
     "fig, axs = plt.subplots(1,2,figsize=(12,5),dpi=100)\n",
     "\n",
-    "loss_train = [x.cpu() for x in model.loss_train]\n",
-    "loss_valid = [x.cpu() for x in model.loss_valid]\n",
+    "loss_train = [x for x in model.logs['loss_train']]\n",
+    "loss_valid = [x for x in model.logs['loss_valid']]\n",
     "\n",
     "# Loss function\n",
     "ax = axs[0]\n",
@@ -514,7 +514,7 @@
     "# Eigenvalues vs epoch\n",
     "ax = axs[1]\n",
     "with torch.no_grad():\n",
-    "    evals_train = np.asarray(torch.cat(model.evals_train).cpu())\n",
+    "    evals_train = np.asarray(model.logs['tica_eigvals'])\n",
     "for i in range(n_eig):\n",
     "    ax.plot(evals_train[:,i],label='Eig. '+str(i+1))\n",
     "ax.set_ylabel('Eigenvalues')\n",