refactor: Simplified if statement to include all fold types

microsoft · Dec 6, 2024 · d2e7d46 · d2e7d46
1 parent bda63b2
commit d2e7d46
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flaml/automl/task/generic_task.py b/flaml/automl/task/generic_task.py
@@ -442,7 +442,7 @@ def prepare_data(
                 X_train_all, y_train_all = shuffle(X_train_all, y_train_all, random_state=RANDOM_SEED)
             if data_is_df:
                 X_train_all.reset_index(drop=True, inplace=True)
-        if split_type in ["group", "uniform", "stratified"] and isinstance(y_train_all, pd.Series):
+        if isinstance(y_train_all, pd.Series):
             y_train_all.reset_index(drop=True, inplace=True)
 
         X_train, y_train = X_train_all, y_train_all