diff --git a/notebooks/Gotthard2/Correction_Gotthard2_NBC.ipynb b/notebooks/Gotthard2/Correction_Gotthard2_NBC.ipynb
index e7c85a868152a875d0b214f5fdaef8c62bd870e0..13ea96c35296ca55fe10524bc3bffca4475715b6 100644
--- a/notebooks/Gotthard2/Correction_Gotthard2_NBC.ipynb
+++ b/notebooks/Gotthard2/Correction_Gotthard2_NBC.ipynb
@@ -443,11 +443,12 @@
     "\n",
     "    # Initialize GH2 data and gain arrays to store in corrected files.\n",
     "    if gh2_detector == \"25um\":\n",
-    "        data_stored = np.zeros((dc[data_sources[0], \"data.adc\"].shape[:2] + (1280 * 2,)), dtype=np.float32)\n",
-    "        gain_stored = np.zeros((dc[data_sources[0], \"data.adc\"].shape[:2] + (1280 * 2,)), dtype=np.uint8)\n",
+    "        dshape = dc[data_sources[0], \"data.adc\"].shape[:2]\n",
+    "        data_stored = np.zeros((dshape + (1280 * 2,)), dtype=np.float32)\n",
+    "        gain_stored = np.zeros((dshape + (1280 * 2,)), dtype=np.uint8)\n",
+    "        mask_stored = np.zeros((dshape + (1280 * 2,)), dtype=np.uint32)\n",
     "    else:\n",
-    "        data_stored = None\n",
-    "        gain_stored = None\n",
+    "        data_stored, gain_stored, mask_stored = None, None, None\n",
     "\n",
     "    for i, (src, mod) in enumerate(zip(data_sources, karabo_da)):\n",
     "        step_timer.start()\n",
@@ -481,9 +482,11 @@
     "        if gh2_detector == \"25um\":\n",
     "            data_stored[..., i::2] = data_corr.copy()\n",
     "            gain_stored[..., i::2] = gain.copy()\n",
+    "            mask_stored[..., i::2] = mask.copy()\n",
     "        else:  # \"50um\"\n",
     "            data_stored = data_corr\n",
     "            gain_stored = gain\n",
+    "            mask_stored = mask\n",
     "\n",
     "    with DataFile(out_file, \"w\") as ofile:\n",
     "        # Create INDEX datasets.\n",