Fix Python typehint with upgraded mypy. (#7513)

dmlc · Dec 16, 2021 · 6f8a463 · 6f8a463
1 parent 70b12d8
commit 6f8a463
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/python-package/xgboost/dask.py b/python-package/xgboost/dask.py
@@ -1859,7 +1859,7 @@ async def _predict_proba_async(
         vstack = update_wrapper(
             partial(da.vstack, allow_unknown_chunksizes=True), da.vstack
         )
-        return _cls_predict_proba(getattr(self, "n_classes_", None), predts, vstack)
+        return _cls_predict_proba(getattr(self, "n_classes_", 0), predts, vstack)
 
     # pylint: disable=missing-function-docstring
     def predict_proba(

diff --git a/python-package/xgboost/sklearn.py b/python-package/xgboost/sklearn.py
@@ -1442,9 +1442,7 @@ def predict_proba(
             iteration_range=iteration_range
         )
         # If model is loaded from a raw booster there's no `n_classes_`
-        return _cls_predict_proba(
-            getattr(self, "n_classes_", None), class_probs, np.vstack
-        )
+        return _cls_predict_proba(getattr(self, "n_classes_", 0), class_probs, np.vstack)
 
 
 @xgboost_model_doc(