correcting pr-auc, roc-auc plots generation (#445)

rudderlabs · Sep 10, 2024 · 84bb7cf · 84bb7cf
1 parent afbde92
commit 84bb7cf
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/predictions/profiles_mlcorelib/trainers/ClassificationTrainer.py b/src/predictions/profiles_mlcorelib/trainers/ClassificationTrainer.py
@@ -191,7 +191,9 @@ def plot_diagnostics(
         label_column: str,
     ) -> None:
         try:
-            predictions = classification_predict_model(model, x)["prediction_score"]
+            predictions = classification_predict_model(model, x, raw_score=True)[
+                "prediction_score_1"
+            ]
 
             y_true = y.to_numpy().reshape(
                 -1,