Fix mypi for multiclass classification for lgbm classifier

nubank · Nov 8, 2023 · c6bbca6 · c6bbca6
1 parent 2730c07
commit c6bbca6
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/fklearn/training/classification.py b/src/fklearn/training/classification.py
@@ -679,6 +679,8 @@ def lgbm_classification_learner(
 
     def p(new_df: pd.DataFrame, apply_shap: bool = False) -> pd.DataFrame:
         predictions = bst.predict(new_df[features].values)
+        if isinstance(predictions, List):
+            predictions = np.ndarray(predictions)
         if is_multiclass_classification:
             col_dict = {prediction_column + "_" + str(key): value
                         for (key, value) in enumerate(predictions.T)}