diff --git a/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb b/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
index 7f4df8a20a170bed676111c2c8d232bb659a506c..7489793dcccde7b06bed6c0362ba4be4e3b1a472 100644
--- a/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
+++ b/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
@@ -141,7 +141,7 @@
     "import matplotlib\n",
     "import matplotlib.pyplot as plt\n",
     "import yaml\n",
-    "from extra_data import RunDirectory, stack_detector_data\n",
+    "from extra_data import by_id, RunDirectory, stack_detector_data\n",
     "from extra_geom import AGIPD_1MGeometry, AGIPD_500K2GGeometry\n",
     "from matplotlib import cm as colormap\n",
     "from matplotlib.colors import LogNorm\n",
@@ -304,7 +304,14 @@
     "\n",
     "    # The trains picked are the unique values of trainTrigger.sequenceStart\n",
     "    # minus the first (previous trigger before this run).\n",
-    "    train_ids = np.unique(seq_start)[1:] + ppu_train_offset\n",
+    "    start_train_ids = np.unique(seq_start)[1:] + ppu_train_offset\n",
+    "\n",
+    "    train_ids = []\n",
+    "    for train_id in start_train_ids:\n",
+    "        n_trains = dc[\n",
+    "            use_ppu_device, 'trainTrigger.numberOfTrains'\n",
+    "        ].select_trains(by_id[[train_id]]).ndarray()[0]\n",
+    "        train_ids.extend(list(range(train_id, train_id + n_trains)))\n",
     "\n",
     "    print(f'PPU device {use_ppu_device} triggered for {len(train_ids)} train(s)')\n",
     "\n",