diff --git a/docs/tutorials/custom_segmentation_trainer.ipynb b/docs/tutorials/custom_segmentation_trainer.ipynb
index 75747fd0bb..fd51a4e45f 100644
--- a/docs/tutorials/custom_segmentation_trainer.ipynb
+++ b/docs/tutorials/custom_segmentation_trainer.ipynb
@@ -73,9 +73,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import os\n",
+    "\n",
     "import lightning\n",
     "import lightning.pytorch as pl\n",
-    "import os\n",
     "import torch\n",
     "from lightning.pytorch.callbacks import ModelCheckpoint\n",
     "from lightning.pytorch.callbacks.callback import Callback\n",