Merge pull request #166 from y0z/feature/fix-ci

Fix for mypy
optuna · Oct 17, 2024 · 403d8e8 · 403d8e8
2 parents 935b449 + 15bdb7a
commit 403d8e8
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/optuna_integration/sklearn/sklearn.py b/optuna_integration/sklearn/sklearn.py
@@ -280,7 +280,10 @@ def _cross_validate_with_pruning(
         if is_classifier(estimator):
             partial_fit_params = self.fit_params.copy()
             y = self.y.values if isinstance(self.y, pd.Series) else self.y
-            classes = np.unique(y)
+            if y is not None:
+                classes = np.unique(y)
+            else:
+                classes = np.array([None])
 
             partial_fit_params.setdefault("classes", classes)