diff --git a/src/cal_tools/agipdlib.py b/src/cal_tools/agipdlib.py
index 121c120c81d12c171c6c0be2679cc1a7ade8415f..3d9f17a7e1d45768e1cd6f95b6f79b0933cc4c8d 100644
--- a/src/cal_tools/agipdlib.py
+++ b/src/cal_tools/agipdlib.py
@@ -1239,21 +1239,21 @@ class AgipdCorrections:
                 # are fixed to the median value.
                 # This is applied for high and medium gain stages
                 for i in range(self.max_cells):
-                    pc_high_m[np.isnan(pc_high_m[..., i])] = pc_high_med[i]
-                    pc_med_m[np.isnan(pc_med_m[..., i])] = pc_med_med[i]
+                    pc_high_m[np.isnan(pc_high_m[..., i]), i] = pc_high_med[i]
+                    pc_med_m[np.isnan(pc_med_m[..., i]), i] = pc_med_med[i]
 
-                    pc_high_l[np.isnan(pc_high_l[..., i])] = pc_high_l_med[i]
-                    pc_med_l[np.isnan(pc_med_l[..., i])] = pc_med_l_med[i]
+                    pc_high_l[np.isnan(pc_high_l[..., i]), i] = pc_high_l_med[i]
+                    pc_med_l[np.isnan(pc_med_l[..., i]), i] = pc_med_l_med[i]
 
                     pc_high_m[(pc_high_m[..., i] < 0.8 * pc_high_med[i]) |
-                              (pc_high_m[..., i] > 1.2 * pc_high_med[i])] = pc_high_med[i]  # noqa
+                              (pc_high_m[..., i] > 1.2 * pc_high_med[i]), i] = pc_high_med[i]  # noqa
                     pc_med_m[(pc_med_m[..., i] < 0.8 * pc_med_med[i]) |
-                             (pc_med_m[..., i] > 1.2 * pc_med_med[i])] = pc_med_med[i]  # noqa
+                             (pc_med_m[..., i] > 1.2 * pc_med_med[i]), i] = pc_med_med[i]  # noqa
 
                     pc_high_l[(pc_high_l[..., i] < 0.8 * pc_high_l_med[i]) |
-                              (pc_high_l[..., i] > 1.2 * pc_high_l_med[i])] = pc_high_l_med[i]  # noqa
+                              (pc_high_l[..., i] > 1.2 * pc_high_l_med[i]), i] = pc_high_l_med[i]  # noqa
                     pc_med_l[(pc_med_l[..., i] < 0.8 * pc_med_l_med[i]) |
-                             (pc_med_l[..., i] > 1.2 * pc_med_l_med[i])] = pc_med_l_med[i]  # noqa
+                             (pc_med_l[..., i] > 1.2 * pc_med_l_med[i]), i] = pc_med_l_med[i]  # noqa
 
                 # ration between HG and MG per pixel per mem cell used
                 # for rel gain calculation