store fold_indices in events during cf.MLEvaluation

haddadanas · Apr 16, 2024 · c8516a5 · c8516a5
1 parent 04db85a
commit c8516a5
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/columnflow/tasks/ml.py b/columnflow/tasks/ml.py
@@ -805,14 +805,14 @@ def run(self):
                 )
 
                 # generate fold indices
-                fold_indices = events.deterministic_seed % self.ml_model_inst.folds
+                events = set_ak_column(events, "fold_indices", events.deterministic_seed % self.ml_model_inst.folds)
 
                 # invoke the optional producer
                 if len(events) and self.preparation_producer_inst:
                     events = self.preparation_producer_inst(
                         events,
                         stats=stats,
-                        fold_indices=fold_indices,
+                        fold_indices=events.fold_indices,
                         ml_model_inst=self.ml_model_inst,
                     )
 
@@ -821,7 +821,7 @@ def run(self):
                     self,
                     events,
                     models,
-                    fold_indices,
+                    events.fold_indices,
                     events_used_in_training=events_used_in_training,
                 )