[AutoTVM] Compatibility improvement with XGBoost v1.3.0 (apache#7076)

neo-ai · Jan 21, 2021 · de52398 · de52398
1 parent c3ca260
commit de52398
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/python/tvm/autotvm/tuner/xgboost_cost_model.py b/python/tvm/autotvm/tuner/xgboost_cost_model.py
@@ -469,7 +469,11 @@ def custom_callback(
     # pylint: disable=import-outside-toplevel
     from xgboost.core import EarlyStopException
     from xgboost.callback import _fmt_metric
-    from xgboost.training import aggcv
+
+    try:
+        from xgboost.training import aggcv
+    except ImportError:
+        from xgboost.callback import _aggcv as aggcv
 
     state = {}
     metric_shortname = metric.split("-")[1]