save n_classes_

csinva · Aug 25, 2023 · 2317acc · 2317acc
1 parent 02f0e9d
commit 2317acc
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/imodels/algebraic/tree_gam.py b/imodels/algebraic/tree_gam.py
@@ -86,6 +86,8 @@ def fit(self, X, y, sample_weight=None):
         X, y = check_X_y(X, y, accept_sparse=False, multi_output=False)
         if isinstance(self, ClassifierMixin):
             check_classification_targets(y)
+            self.classes_, y = np.unique(y, return_inverse=True)
+
         sample_weight = _check_sample_weight(sample_weight, X, dtype=None)
 
         # split into train and validation for early stopping