diff --git a/CHANGELOG.md b/CHANGELOG.md
index 95afd0f0..21694246 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,7 @@
 ### 0.2.6 (UNRELEASED)
 
 ### Datasets
+* Correct stage-based conditions mentioned in notebook tutorials [#92](https://github.com/a-r-j/ProteinWorkshop/pull/92)
 * Add stage-based conditions to `setup` in `ProteinDataModule` [#72](https://github.com/a-r-j/ProteinWorkshop/pull/72)
 * Improves support for datamodules with multiple test sets. Generalises this to support GO and FOLD. Also adds multiple seq ID.-based splits for GO. [#72](https://github.com/a-r-j/ProteinWorkshop/pull/72)
 * Add redownload checks for already downloaded datasets and harmonise pdb download interface [#86](https://github.com/a-r-j/ProteinWorkshop/pull/86)
diff --git a/notebooks/adding_new_dataset_tutorial.ipynb b/notebooks/adding_new_dataset_tutorial.ipynb
index 24bfcb9b..8d84f207 100644
--- a/notebooks/adding_new_dataset_tutorial.ipynb
+++ b/notebooks/adding_new_dataset_tutorial.ipynb
@@ -352,7 +352,7 @@
     "cfg = config.validate_config(cfg)\n",
     "\n",
     "datamodule = hydra.utils.instantiate(cfg.dataset.datamodule)\n",
-    "datamodule.setup(\"train\")\n",
+    "datamodule.setup(\"fit\")\n",
     "dl = datamodule.train_dataloader()\n",
     "\n",
     "for i in dl:\n",
@@ -448,7 +448,7 @@
     "cfg = config.validate_config(cfg)\n",
     "\n",
     "datamodule = hydra.utils.instantiate(cfg)\n",
-    "datamodule.setup(\"train\")\n",
+    "datamodule.setup(\"fit\")\n",
     "dl = datamodule.train_dataloader()\n",
     "\n",
     "for i in dl:\n",
diff --git a/notebooks/customizing_existing_dataset_tutorial.ipynb b/notebooks/customizing_existing_dataset_tutorial.ipynb
index 165ae867..c8dc43d3 100644
--- a/notebooks/customizing_existing_dataset_tutorial.ipynb
+++ b/notebooks/customizing_existing_dataset_tutorial.ipynb
@@ -118,7 +118,7 @@
     "cfg = config.validate_config(cfg)\n",
     "\n",
     "datamodule = hydra.utils.instantiate(cfg.dataset.datamodule)\n",
-    "datamodule.setup(\"train\")\n",
+    "datamodule.setup(\"fit\")\n",
     "dl = datamodule.train_dataloader()\n",
     "\n",
     "for i in dl:\n",
@@ -207,7 +207,7 @@
     "cfg = config.validate_config(cfg)\n",
     "\n",
     "datamodule = hydra.utils.instantiate(cfg)\n",
-    "datamodule.setup(\"train\")\n",
+    "datamodule.setup(\"fit\")\n",
     "dl = datamodule.train_dataloader()\n",
     "\n",
     "for i in dl:\n",