diff --git a/docs/notebooks/segmentation_pipeline.ipynb b/docs/notebooks/segmentation_pipeline.ipynb
index 326ab9d78dabbea918af8e1662965c373c3716c4..13c986e90beee4dc0dfe82c18a3e076b68882949 100644
--- a/docs/notebooks/segmentation_pipeline.ipynb
+++ b/docs/notebooks/segmentation_pipeline.ipynb
@@ -528,7 +528,7 @@
    "outputs": [],
    "source": [
     "# defining model\n",
-    "model = qim3d.ml.models.UNet(size = 'small', dropout = 0.25)"
+    "model = qim3d.ml.models.UNet2D(size = 'small', dropout = 0.25)"
    ]
   },
   {