diff --git a/cal_tools/cal_tools/agipdlib.py b/cal_tools/cal_tools/agipdlib.py
index 5451d5b9316d72fd53a358a48ba6a7049bc363ce..18328ff08982653d5b43339b60c2ab9154013a2f 100644
--- a/cal_tools/cal_tools/agipdlib.py
+++ b/cal_tools/cal_tools/agipdlib.py
@@ -371,17 +371,16 @@ class AgipdCorrections:
         dark_max = self.cm_dark_max
         fraction = self.cm_dark_fraction
         n_itr = self.cm_n_itr
-        
-        cell_id = self.shared_dict[i_proc]['cellId']
-        train_id = self.shared_dict[i_proc]['trainId']
         n_img = self.shared_dict[i_proc]['nImg'][0]
+        cell_id = self.shared_dict[i_proc]['cellId'][:n_img]
+        train_id = self.shared_dict[i_proc]['trainId'][:n_img]
         cell_ids = cell_id[train_id == train_id[0]]
         n_cells = cell_ids.size
         data = self.shared_dict[i_proc]['data'][:n_img].reshape(-1, n_cells,
                                                                 8, 64, 2, 64)
 
         # Loop over iterations
-        for i in range(n_itr):
+        for _ in range(n_itr):
             # Loop over rows of cells
             first = 0
             for cell_row in range(11):