diff --git a/docs/notebooks/dwi_gp_estimation.ipynb b/docs/notebooks/dwi_gp_estimation.ipynb
index 2cce4b3e..a26dbcb4 100644
--- a/docs/notebooks/dwi_gp_estimation.ipynb
+++ b/docs/notebooks/dwi_gp_estimation.ipynb
@@ -29,23 +29,9 @@
     "from dipy.core.gradients import get_bval_indices\n",
     "from dipy.io import read_bvals_bvecs\n",
     "from dipy.segment.mask import median_otsu\n",
-    "\n",
     "from scipy.ndimage import binary_dilation\n",
     "from skimage.morphology import ball\n",
     "\n",
-    "from nireports.reportlets.modality.dwi import plot_dwi"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "id": "69a3bc6b4fbe7036",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "seed = 1234\n",
-    "rng = np.random.default_rng(seed)\n",
-    "\n",
     "name = \"sherbrooke_3shell\"\n",
     "\n",
     "dwi_fname, bval_fname, bvec_fname = dpd.get_fnames(name=name)\n",
@@ -131,6 +117,8 @@
     }
    ],
    "source": [
+    "from nireports.reportlets.modality.dwi import plot_dwi\n",
+    "\n",
     "affine = nib.load(dwi_fname).affine\n",
     "plot_dwi(shell_data[..., dwi_vol_idx], affine, gradient=np.concatenate((np.squeeze(X_test), [1000])), output_file=\"sherbrooke_3shell_b1k_data.svg\")"
    ]
@@ -221,7 +209,14 @@
      "output_type": "display_data"
     }
    ],
-   "source": "plot_dwi(dwi_sim, affine, gradient=np.concatenate((np.squeeze(X_test), [1000])), output_file=\"sherbrooke_3shell_b1k_gp_noopt_pred.svg\")"
+   "source": [
+    "plot_dwi(\n",
+    "    dwi_sim,\n",
+    "    affine,\n",
+    "    gradient=np.concatenate((np.squeeze(X_test), [1000])),\n",
+    "    output_file=\"sherbrooke_3shell_b1k_gp_noopt_pred.svg\",\n",
+    ")"
+   ]
   },
   {
    "cell_type": "code",
@@ -449,8 +444,15 @@
    ],
    "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])), output_file=\"sherbrooke_3shell_b1k_gp_opt_pred.svg\")"
+    "plot_dwi(\n",
+    "    shell_data[..., dwi_vol_idx], affine, gradient=np.concatenate((np.squeeze(X_test), [1000]))\n",
+    ")\n",
+    "plot_dwi(\n",
+    "    dwi_sim2,\n",
+    "    affine,\n",
+    "    gradient=np.concatenate((np.squeeze(X_test), [1000])),\n",
+    "    output_file=\"sherbrooke_3shell_b1k_gp_opt_pred.svg\",\n",
+    ")"
    ]
   }
  ],