diff --git a/notebooks/Chapter04/01-Setting up Experiment Harness.ipynb b/notebooks/Chapter04/01-Setting up Experiment Harness.ipynb
index adae24b..c923b7c 100644
--- a/notebooks/Chapter04/01-Setting up Experiment Harness.ipynb	
+++ b/notebooks/Chapter04/01-Setting up Experiment Harness.ipynb	
@@ -802,8 +802,13 @@
     }
    ],
    "source": [
-    "test_mask = (exp_block_df.timestamp.dt.year==2014) & (exp_block_df.timestamp.dt.month==2)\n",
-    "val_mask = (exp_block_df.timestamp.dt.year==2014) & (exp_block_df.timestamp.dt.month==1)\n",
+    "# test_mask = (exp_block_df.timestamp.dt.year==2014) & (exp_block_df.timestamp.dt.month==2)\n",
+    "# val_mask = (exp_block_df.timestamp.dt.year==2014) & (exp_block_df.timestamp.dt.month==1)\n",
+    "\n",
+    "\"\"\"To get the last year and last months values from the dataset without specifying manually.\"\"\"\n",
+    "\n",
+    "test_mask = (exp_block_df.timestamp.dt.year==exp_block_df.timestamp.dt.to_period('M').unique()[-1].year) & (exp_block_df.timestamp.dt.month==exp_block_df.timestamp.dt.to_period('M').unique()[-1].month)\n",
+    "val_mask = (exp_block_df.timestamp.dt.year==exp_block_df.timestamp.dt.to_period('M').unique()[-1].year) & (exp_block_df.timestamp.dt.month==exp_block_df.timestamp.dt.to_period('M').unique()[-2].month)\n",
     "\n",
     "train = exp_block_df[~(val_mask|test_mask)]\n",
     "val = exp_block_df[val_mask]\n",
@@ -1188,8 +1193,11 @@
     }
    ],
    "source": [
-    "test_mask = (exp_block_df.timestamp.dt.year==2014) & (exp_block_df.timestamp.dt.month==2)\n",
-    "val_mask = (exp_block_df.timestamp.dt.year==2014) & (exp_block_df.timestamp.dt.month==1)\n",
+    "# test_mask = (exp_block_df.timestamp.dt.year==2014) & (exp_block_df.timestamp.dt.month==2)\n",
+    "# val_mask = (exp_block_df.timestamp.dt.year==2014) & (exp_block_df.timestamp.dt.month==1)\n",
+    "\n",
+    "test_mask = (exp_block_df.timestamp.dt.year==exp_block_df.timestamp.dt.to_period('M').unique()[-1].year) & (exp_block_df.timestamp.dt.month==exp_block_df.timestamp.dt.to_period('M').unique()[-1].month)\n",
+    "val_mask = (exp_block_df.timestamp.dt.year==exp_block_df.timestamp.dt.to_period('M').unique()[-1].year) & (exp_block_df.timestamp.dt.month==exp_block_df.timestamp.dt.to_period('M').unique()[-2].month)\n",
     "\n",
     "train = exp_block_df[~(val_mask|test_mask)]\n",
     "val = exp_block_df[val_mask]\n",
@@ -1227,7 +1235,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.0"
+   "version": "3.9.18"
   }
  },
  "nbformat": 4,