fix lint

apache · Sep 14, 2022 · 918ea89 · 918ea89
1 parent 79599af
commit 918ea89
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/python/tvm/meta_schedule/cost_model/xgb_model.py b/python/tvm/meta_schedule/cost_model/xgb_model.py
@@ -36,7 +36,8 @@
 from .metric import max_curve
 
 
-def optional_xgboost_callback(XGBoostCustomCallback):
+def optional_xgboost_callback(cls):
+    """Decorator for importing TraningCallback from xgboost"""
     # pylint:disable = import-outside-toplevel
     try:
         from xgboost.callback import TrainingCallback  # type: ignore
@@ -46,7 +47,7 @@ def optional_xgboost_callback(XGBoostCustomCallback):
         class TrainingCallback:  # type: ignore
             pass
 
-    class OptXGBoostCustomCallback(XGBoostCustomCallback, TrainingCallback):  # type: ignore
+    class OptXGBoostCustomCallback(cls, TrainingCallback):  # type: ignore
         pass
 
     return OptXGBoostCustomCallback