diff --git a/tests/runtime/test_xgboost.py b/tests/runtime/test_xgboost.py index 25b8eaaf1f..7d524c2892 100644 --- a/tests/runtime/test_xgboost.py +++ b/tests/runtime/test_xgboost.py @@ -2,7 +2,7 @@ import pytest dask_ml = pytest.importorskip("dask_ml") -xgb = pytest.importorskip("xgboost") +dxgb = pytest.importorskip("xgboost.dask") def test_xgboost_distributed_training(small_client): @@ -20,8 +20,8 @@ def test_xgboost_distributed_training(small_client): ) # Create the XGBoost DMatrix for our training and testing splits - dtrain = xgb.dask.DaskDMatrix(small_client, X_train, y_train) - dtest = xgb.dask.DaskDMatrix(small_client, X_test, y_test) + dtrain = dxgb.DaskDMatrix(small_client, X_train, y_train) + dtest = dxgb.DaskDMatrix(small_client, X_test, y_test) # Set model parameters (XGBoost defaults) params = { @@ -32,10 +32,10 @@ def test_xgboost_distributed_training(small_client): "objective": "reg:squarederror", "grow_policy": "depthwise", } - output = xgb.dask.train( + output = dxgb.train( small_client, params, dtrain, num_boost_round=5, evals=[(dtrain, "train")] ) # make predictions - y_pred = xgb.dask.predict(small_client, output, dtest) + y_pred = dxgb.predict(small_client, output, dtest) assert y_pred.shape[0] == y_test.shape[0].compute()