diff --git a/code/bolsonaro/trainer.py b/code/bolsonaro/trainer.py
index c2bd7673262bb5bfd18ae617b1ba0c62b1092506..9d4911ea3704381d4b101a74267c760886f68d5d 100644
--- a/code/bolsonaro/trainer.py
+++ b/code/bolsonaro/trainer.py
@@ -145,16 +145,14 @@ class Trainer(object):
         from sklearn.preprocessing import normalize
         import itertools
 
-        if type(model) in [OmpForestRegressor, SimilarityForestRegressor, KMeansForestRegressor, EnsembleSelectionForestRegressor]:
-            estimators = model.forest
-        elif type(model) in [OmpForestBinaryClassifier, OmpForestMulticlassClassifier]:
+        if type(model) in [OmpForestRegressor, SimilarityForestRegressor, KMeansForestRegressor, EnsembleSelectionForestRegressor,
+            OmpForestBinaryClassifier, OmpForestMulticlassClassifier]:
             estimators = model.forest
+            estimators = np.asarray(estimators)[model._omp.coef_ != 0]
         elif type(model) in [RandomForestRegressor, RandomForestClassifier]:
             estimators = model.estimators_
 
-        predictions = list()
-        for ti in estimators:
-            predictions.append(ti.predict(X))
+        predictions = np.array([tree.predict(X) for tree in estimators])
 
         predictions = normalize(predictions)