diff --git a/notebooks/LPDMini/LPD_Mini_Correct.ipynb b/notebooks/LPDMini/LPD_Mini_Correct.ipynb
index c541296565e244793b513c785b892eb0a57f9f13..34e6214e3457475069da145abac899168ca1810f 100644
--- a/notebooks/LPDMini/LPD_Mini_Correct.ipynb
+++ b/notebooks/LPDMini/LPD_Mini_Correct.ipynb
@@ -436,6 +436,8 @@
     "    in_sliced = np.zeros(data_shape, dtype=in_raw.dtype)\n",
     "    for i, sl in enumerate(raw_data_slices):\n",
     "        in_sliced[:, i*32:(i+1)*32] = in_raw[..., sl, :]\n",
+    "        \n",
+    "    output_shape = (data_shape[0], nmods, 32, 256)\n",
     "    \n",
     "    # Allocate output arrays.\n",
     "    out_data = np.zeros(in_sliced.shape, dtype=np.float32)\n",
@@ -464,11 +466,10 @@
     "                                   chunks=(min(chunks_ids, in_cell.shape[0]),))\n",
     "            outp_source.create_key('image.pulseId', data=in_pulse,\n",
     "                                   chunks=(min(chunks_ids, in_pulse.shape[0]),))\n",
-    "            outp_source.create_key('image.data',\n",
-    "                                   data=out_data.reshape(data_shape[0], nmods, 32, 256),\n",
-    "                                   chunks=(min(chunks_data, out_data.shape[0]), nmods, 32, 256))\n",
-    "            outp_source.create_compressed_key('image.gain', data=out_gain)\n",
-    "            outp_source.create_compressed_key('image.mask', data=out_mask)\n",
+    "            outp_source.create_key('image.data', data=out_data.reshape(output_shape),\n",
+    "                                   chunks=(min(chunks_data, out_data.shape[0]), 1, 32, 256))\n",
+    "            outp_source.create_compressed_key('image.gain', data=out_gain.reshape(output_shape))\n",
+    "            outp_source.create_compressed_key('image.mask', data=out_mask.reshape(output_shape))\n",
     "\n",
     "    write_time = perf_counter() - start\n",
     "    \n",
@@ -509,6 +510,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# This geometry is arbitrary, we just want to show all the modules\n",
     "geom = xg.LPD_MiniGeometry.from_module_positions(\n",
     "    [(0, i * 40) for i in range(nmods)]\n",
     ")\n",
@@ -598,7 +600,7 @@
    "outputs": [],
    "source": [
     "fig, ax = plt.subplots(num=3, figsize=(15, 15), clear=True, nrows=1, ncols=1)\n",
-    "geom.plot_data_fast(data.mean(axis=1), ax=ax, vmin=vmin, vmax=vmax)\n",
+    "geom.plot_data_fast(data.mean(axis=0), ax=ax, vmin=vmin, vmax=vmax)\n",
     "pass"
    ]
   },