Merge pull request #325 from NeroHin/master

feat: Add estimator type with reg and clf for NGBClassifier and NGBRegressor
stanfordmlgroup · Oct 29, 2024 · acf6b27 · acf6b27
2 parents 38c2245 + ce588f0
commit acf6b27
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/ngboost/api.py b/ngboost/api.py
@@ -97,6 +97,8 @@ def __init__(
             early_stopping_rounds,
         )
 
+        self._estimator_type = "regressor"
+
     def __getstate__(self):
         state = super().__getstate__()
         # Remove the unpicklable entries.
@@ -172,6 +174,7 @@ def __init__(
             tol,
             random_state,
         )
+        self._estimator_type = "classifier"
 
     def predict_proba(self, X, max_iter=None):
         """