diff --git a/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb b/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
index 9c38bd80c4b8e4713866bdfca41d78fe404e48a0..cbb323604ffc01f73d5be95e2a3044562f530ef2 100644
--- a/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
+++ b/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
@@ -882,24 +882,15 @@
     "    \"\"\"\n",
     "    run_data = RunDirectory(data_folder, include)\n",
     "    if tid is not None:\n",
-    "        tid, data = run_data.select(f'{detector_id}/DET/*', source).train_from_id(tid)\n",
+    "        tid, data = run_data.select(\n",
+    "            f'{detector_id}/DET/*', source).train_from_id(tid, keep_dims=True)\n",
     "    else:\n",
-    "        tid, data = next(iter(run_data.select(f'{detector_id}/DET/*', source).trains(require_all=True)))\n",
-    "        \n",
-    "    # TODO: remove and use the keep_dims version after updating Extra-data.\n",
-    "    # Avoid using default axis with sources of an expected scalar value per train.\n",
-    "    nfrm = cell_sel.get_cells_on_trains([tid]).sum()\n",
-    "    if nfrm == 1 and source in ['image.blShift', 'image.cellId', 'image.pulseId']:\n",
-    "        axis = 0\n",
-    "    else:\n",
-    "        axis = -3\n",
+    "        # A first full trainId for all available modules is of interest.\n",
+    "        tid, data = next(run_data.select(\n",
+    "            f'{detector_id}/DET/*', source).trains(require_all=True, keep_dims=True))\n",
     "\n",
     "    stacked_data = stack_detector_data(\n",
-    "        train=data, data=source, fillvalue=fillvalue, modules=modules, axis=axis)\n",
-    "    # Add cellId dimension when correcting one cellId only.\n",
-    "    # avoid adding pulse dims for raw data.\n",
-    "    if (nfrm == 1 and data_folder != run_folder):\n",
-    "        stacked_data = stacked_data[np.newaxis, ...]\n",
+    "        train=data, data=source, fillvalue=fillvalue, modules=modules)\n",
     "\n",
     "    return tid, stacked_data"
    ]