minor fix on compatibility (#454)

Co-authored-by: Qingyun Wu <qxw5138@psu.edu>
openml · Mar 1, 2022 · db5c550 · db5c550
1 parent 548b56e
commit db5c550
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/frameworks/flaml/exec.py b/frameworks/flaml/exec.py
@@ -53,7 +53,9 @@ def run(dataset, config):
     with Timer() as predict:
         predictions = aml.predict(X_test)
     probabilities = aml.predict_proba(X_test) if is_classification else None
-    labels = aml.classes_ if is_classification else None
+    labels = None
+    if is_classification:
+        labels = aml.classes_ if isinstance(aml.classes_, list) else aml.classes_.tolist()
     return result(  
                     output_file=config.output_predictions_file,
                     probabilities=probabilities,