diff --git a/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb b/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
index 9c38bd80c4b8e4713866bdfca41d78fe404e48a0..679a3172d2c4263974ae3ecfd893fe5e4cdf1019 100644
--- a/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
+++ b/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
@@ -882,7 +882,8 @@
     "    \"\"\"\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",