diff --git a/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb b/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb
index 144d8ecb9d54d1527cb84e49452725279eb1d1fc..8f91e9263fbb1efd967861462b289bddefc3314d 100644
--- a/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb
+++ b/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb
@@ -867,7 +867,7 @@
     "display(Markdown((f\"#### The per pixel maximum of train {tid} of the GAIN data\")))\n",
     "\n",
     "fig, ax = plt.subplots(figsize=(18, 10))\n",
-    "gain_max = np.max(gain_train_cells, axis=1)\n",
+    "gain_max = np.max(gain_train_cells, axis=(1, 2))\n",
     "geom.plot_data_fast(\n",
     "    gain_max,\n",
     "    ax=ax,\n",