diff --git a/code/bolsonaro/models/omp_forest.py b/code/bolsonaro/models/omp_forest.py
index 5b947d327693020b51c7da778d4855274454de93..35c2f8fa0e46e2c87a76be438a9968665e7f5b29 100644
--- a/code/bolsonaro/models/omp_forest.py
+++ b/code/bolsonaro/models/omp_forest.py
@@ -136,14 +136,11 @@ class SingleOmpForest(OmpForest):
         :param X: a Forest
         :return: a np.array of the predictions of the entire forest
         """
-        forest_predictions = self._base_estimator_predictions(X).T
+        forest_predictions = np.array([tree.predict(X) for tree in self._base_forest_estimator.estimators_])
 
         if self._models_parameters.normalize_D:
             forest_predictions /= self._forest_norms
 
         weights = self._omp.coef_
-        omp_trees_indices = np.nonzero(weights)[0]
-
-        select_trees = np.mean(forest_predictions[omp_trees_indices], axis=0)
-        print(len(omp_trees_indices))
+        select_trees = np.mean(forest_predictions[weights != 0], axis=0)
         return select_trees
diff --git a/code/bolsonaro/models/omp_forest_classifier.py b/code/bolsonaro/models/omp_forest_classifier.py
index a51405a6a3278bb86dd52d011b599175bbfc7482..3051fad09c04c34bab5f7035f7392c09313c7b30 100644
--- a/code/bolsonaro/models/omp_forest_classifier.py
+++ b/code/bolsonaro/models/omp_forest_classifier.py
@@ -40,9 +40,7 @@ class OmpForestBinaryClassifier(SingleOmpForest):
             forest_predictions /= self._forest_norms
 
         weights = self._omp.coef_
-        omp_trees_indices = np.nonzero(weights)
-
-        omp_trees_predictions = forest_predictions[omp_trees_indices].T[1]
+        omp_trees_predictions = forest_predictions[weights != 0].T[1]
 
         # Here forest_pred is the probability of being class 1.