diff --git a/code/bolsonaro/trainer.py b/code/bolsonaro/trainer.py index 21b7ab13463ba4bca91de05675003a5541133296..1d06dbe7291a41b8dc50f1ca7d8b719f0c197512 100644 --- a/code/bolsonaro/trainer.py +++ b/code/bolsonaro/trainer.py @@ -155,7 +155,7 @@ class Trainer(object): def _evaluate_predictions(self, model, X, aggregation_function): if type(model) in [OmpForestRegressor, SimilarityForestRegressor, KMeansForestRegressor, EnsembleSelectionForestRegressor, - OmpForestBinaryClassifier, OmpForestMulticlassClassifier]: + OmpForestBinaryClassifier, OmpForestMulticlassClassifier, SimilarityForestClassifier, KMeansForestClassifier, EnsembleSelectionForestClassifier]: estimators = model.forest estimators = np.asarray(estimators)[model._omp.coef_ != 0] elif type(model) in [RandomForestRegressor, RandomForestClassifier]: