diff --git a/python-package/xgboost/dask.py b/python-package/xgboost/dask.py index 665f7c264..cb972afa3 100644 --- a/python-package/xgboost/dask.py +++ b/python-package/xgboost/dask.py @@ -1859,7 +1859,7 @@ class DaskXGBClassifier(DaskScikitLearnBase, XGBClassifierBase): 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 index 07206bb55..949dae7b4 100644 --- a/python-package/xgboost/sklearn.py +++ b/python-package/xgboost/sklearn.py @@ -1442,9 +1442,7 @@ class XGBClassifier(XGBModel, XGBClassifierBase): 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(