diff --git a/docs/notebooks/dwi_gp_estimation.ipynb b/docs/notebooks/dwi_gp_estimation.ipynb
index 777f288d..b4e01e00 100644
--- a/docs/notebooks/dwi_gp_estimation.ipynb
+++ b/docs/notebooks/dwi_gp_estimation.ipynb
@@ -132,7 +132,7 @@
    ],
    "source": [
     "affine = nib.load(dwi_fname).affine\n",
-    "plot_dwi(shell_data[..., dwi_vol_idx], affine, gradient=np.concatenate((np.squeeze(X_test), [1000])))"
+    "plot_dwi(shell_data[..., dwi_vol_idx], affine, gradient=np.concatenate((np.squeeze(X_test), [1000])), output_file=\"sherbrooke_3shell_b1k_data.svg\")"
    ]
   },
   {
@@ -221,9 +221,7 @@
      "output_type": "display_data"
     }
    ],
-   "source": [
-    "plot_dwi(dwi_sim, affine, gradient=np.concatenate((np.squeeze(X_test), [1000])))"
-   ]
+   "source": "plot_dwi(dwi_sim, affine, gradient=np.concatenate((np.squeeze(X_test), [1000])), output_file=\"sherbrooke_3shell_b1k_gp_noopt_pred.svg\")"
   },
   {
    "cell_type": "code",
@@ -452,7 +450,7 @@
    "source": [
     "plot_dwi(dwi_sim, affine, gradient=np.concatenate((np.squeeze(X_test), [1000])))\n",
     "plot_dwi(shell_data[..., dwi_vol_idx], affine, gradient=np.concatenate((np.squeeze(X_test), [1000])))\n",
-    "plot_dwi(dwi_sim2, affine, gradient=np.concatenate((np.squeeze(X_test), [1000])))"
+    "plot_dwi(dwi_sim2, affine, gradient=np.concatenate((np.squeeze(X_test), [1000])), output_file=\"sherbrooke_3shell_b1k_gp_opt_pred.svg\")"
    ]
   },
   {