Updated finding max AUC of classifier ensembles

varunlakshmanan · Jun 8, 2020 · 9dcd156 · 9dcd156
1 parent 8c9ee8c
commit 9dcd156
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/ensemble_models.py b/src/ensemble_models.py
@@ -36,10 +36,11 @@ def ensemble_models(optimized_estimators, x_train, y_train, x_test, y_test, is_c
             voting_estimator.fit(x_train, y_train)
             y_predictions = voting_estimator.predict(x_test)
             auc = roc_auc_score(y_test, y_predictions)
+
             if auc > max_auc:
                 best_voting_estimator = voting_estimator
+                max_auc = auc
 
-        print("Highest AUC: " + str(max_auc))
         return best_voting_estimator
 
     else: