diff --git a/nnfwtbn/model.py b/nnfwtbn/model.py index 98982a6b248ed2e1626082fe89656cf00de6e462..5a35d21f01e5c1ff75b01c71052dc498ed00ad66 100644 --- a/nnfwtbn/model.py +++ b/nnfwtbn/model.py @@ -361,7 +361,7 @@ class MixedCV(CrossValidator): def select_training_slices(self, fold_i, for_predicting = False): """ - Returns array of arrays with integers corresponding + Returns array with integers corresponding to the data slices used in training fold_i. If 'for_predicting' is set to True only one slice is returned for each fold so that the folds are equally represented @@ -390,7 +390,7 @@ class MixedCV(CrossValidator): for sl in slices: if sl not in exclusive_slices and sl in uniq_el(all_slices_for_folds[i:]): exclusive_slices.append(sl) - return exclusive_slices[fold_i] + return [exclusive_slices[fold_i]] def select_training(self, df, fold_i, for_predicting = False): """ @@ -399,7 +399,7 @@ class MixedCV(CrossValidator): """ selected = np.zeros(len(df), dtype='bool') slices = self.select_training_slices(fold_i, for_predicting = for_predicting) - for slice_i in np.array(slices).flatten(): + for slice_i in slices: selected = selected | self.select_slice(df, slice_i) return selected