diff --git a/src/exdf/data_reduction/builtins.py b/src/exdf/data_reduction/builtins.py
index b35f48540732a54e05cc6c8c17249eb69e867c3f..88400cfd13f1b40c21378e35b718ee56d713759b 100644
--- a/src/exdf/data_reduction/builtins.py
+++ b/src/exdf/data_reduction/builtins.py
@@ -344,13 +344,9 @@ class LitFrames(ReductionMethod):
                  train_sel, nfrm, cell_id, [sel_type])
 
             # select frames for every train
-            i0 = 0
-            for i in range(len(train_sel)):
-                iN = i0 + nfrm[i]
-                self.select_xtdf(src,
-                                 by_id[[train_sel[i]]],
-                                 frame_sel[i0:iN])
-                i0 = iN
+            train_masks = np.split(frame_sel, np.cumsum(nfrm))
+            for train_id, mask in zip(train_sel, train_masks):
+                self.select_xtdf(src, by_id[[train_id]], mask)
 
     def print_report(self, sel, max_lines=20):
         report = sel.report()