diff --git a/notebooks/AGIPD/Characterize_AGIPD_Gain_Darks_NBC.ipynb b/notebooks/AGIPD/Characterize_AGIPD_Gain_Darks_NBC.ipynb
index 32068cbac86097dcbff72cd9a289cf202289eb7a..f3998ba558eb853de3d3655a2b095362be97ddcd 100644
--- a/notebooks/AGIPD/Characterize_AGIPD_Gain_Darks_NBC.ipynb
+++ b/notebooks/AGIPD/Characterize_AGIPD_Gain_Darks_NBC.ipynb
@@ -222,7 +222,8 @@
     "    # validate that there are trains and that data sources are\n",
     "    # present for any of the selected modules.\n",
     "    n_trains = len(dc.train_ids)\n",
-    "\n",
+    "    if n_trains < min_trains:\n",
+    "        print(f\"WARNING: {dc.files} have less than minimum trains: {min_trains}.\")\n",
     "    if n_trains == 0:\n",
     "        raise ValueError(f\"No images to process for run: {run_dict['number']}\")\n",
     "\n",
@@ -436,30 +437,6 @@
     "    ]"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Check if max_trains can be processed.\n",
-    "\n",
-    "# more relevant if running on multiple modules (i.e. within notebook)\n",
-    "# mem_cells * gains * n_constants * modules * agipd_[x,y]image_size * 2\n",
-    "av_mem = psutil.virtual_memory().available\n",
-    "possible_trains = av_mem // (352 * 3 * 3 * len(modules) * 131072 * 2)\n",
-    "if max_trains == 0:\n",
-    "    max_trains = max(max_trains_list)\n",
-    "if max_trains > possible_trains:\n",
-    "    max_trains = possible_trains\n",
-    "    print(\n",
-    "        f\"WARNING: available memory for processing is { av_mem / 1e9:.02f} GB.\"\n",
-    "        f\" Modifing max_trains to process to {max_trains}\")\n",
-    "\n",
-    "for run_dict in runs_dict.values():\n",
-    "    run_dict[\"dc\"] = run_dict[\"dc\"].select_trains(np.s_[:max_trains])"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -475,7 +452,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "parallel_num_procs = min(12, len(modules)*3)\n",
+    "parallel_num_procs = min(6, len(modules)*3)\n",
     "parallel_num_threads = multiprocessing.cpu_count() // parallel_num_procs\n",
     "print(f\"Will use {parallel_num_procs} processes with {parallel_num_threads} threads each\")\n",
     "\n",
@@ -486,14 +463,11 @@
     "    # Select the corresponding module channel.\n",
     "    instrument_src_mod = instrument_src.format(channel)\n",
     "\n",
-    "    run_dc = runs_dict[\"dc\"]\n",
+    "    run_dc = runs_dict[\"dc\"].select(instrument_src_mod, require_all=True)\n",
+    "    if max_trains != 0:\n",
+    "        run_dc = run_dc.select_trains(np.s_[:max_trains])\n",
     "    gain_index = runs_dict[\"gain\"]\n",
     "\n",
-    "    if run_dc[instrument_src_mod, \"image.data\"].shape[0] < min_trains:\n",
-    "        print(\n",
-    "            f\"WARNING: {run_dc.files} have less than \"\n",
-    "            \"minimum trains: {min_trains}.\")\n",
-    "\n",
     "    # Read module's image and cellId data.\n",
     "    im = run_dc[instrument_src_mod, \"image.data\"].ndarray()\n",
     "    cell_ids = np.squeeze(run_dc[instrument_src_mod, \"image.cellId\"].ndarray())\n",