diff --git a/code/bolsonaro/models/omp_forest_regressor.py b/code/bolsonaro/models/omp_forest_regressor.py
index 3e95d12b335e1cca9645517c0dfcc08ff8b73e5a..b4896cda87396da32c37eccdbe7cf8c4f743976e 100644
--- a/code/bolsonaro/models/omp_forest_regressor.py
+++ b/code/bolsonaro/models/omp_forest_regressor.py
@@ -38,7 +38,8 @@ class OmpForestRegressor(BaseEstimator):
 
     def fit(self, X_forest, y_forest, X_omp, y_omp):
         self._forest = self._train_forest(X_forest, y_forest)
-        self._weights = self._extract_subforest(X_omp, y_omp)
+        self._omp = self._extract_subforest(X_omp, y_omp)
+        self._weights = self._omp.coef_
         return self
 
     def score_regressor(self, X, y):
@@ -56,8 +57,8 @@ class OmpForestRegressor(BaseEstimator):
         if self._models_parameters.normalize_D:
             D /= self._forest_norms
 
-        # TODO: use self._models_parameters.normalize_weights here
-        predictions = D @ self._weights
+        # TODO: use self._models_parameters.normalize_weights here?
+        predictions = self._omp.predict(D)
 
         return predictions
 
@@ -109,10 +110,7 @@ class OmpForestRegressor(BaseEstimator):
             fit_intercept=False, normalize=False)
         self._logger.debug("Apply orthogonal maching pursuit on forest for {} extracted trees."
                            .format(self._models_parameters.extracted_forest_size))
-        omp.fit(D, y)
-        weights = omp.coef_
-        # question: why not to use directly the omp estimator instead of bypassing it using the coefs?
-        return weights
+        return omp.fit(D, y)
 
     def _forest_prediction(self, X):
         return np.array([tree.predict(X) for tree in self._forest]).T