Fix xgboost pruning callback (#4921)

* allow XGBoostPruningCallback for xgboost >= 2.0.0 * Run black --------- Co-authored-by: Yiwen Yuan <yiwen@kumo.ai>
optuna · Sep 12, 2023 · 9494173 · 9494173
1 parent a30bc0b
commit 9494173
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/optuna/integration/xgboost.py b/optuna/integration/xgboost.py
@@ -11,7 +11,9 @@
     xgboost_version = xgb.__version__.split(".")
     xgboost_major_version = int(xgboost_version[0])
     xgboost_minor_version = int(xgboost_version[1])
-    use_callback_cls = xgboost_major_version >= 1 and xgboost_minor_version >= 3
+    use_callback_cls = (
+        xgboost_major_version >= 1 and xgboost_minor_version >= 3
+    ) or xgboost_major_version >= 2
 
 _doc = """Callback for XGBoost to prune unpromising trials.