Persist data in dask test. (#7077)

This commit is contained in:
Jiaming Yuan 2021-07-06 11:47:17 +08:00 committed by GitHub
parent b56d3d5d5c
commit ffa66aace0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -184,6 +184,9 @@ def test_dask_predict_shape_infer(client: "Client") -> None:
def run_boost_from_prediction(
X: xgb.dask._DaskCollection, y: xgb.dask._DaskCollection, tree_method: str, client: "Client"
) -> None:
X = client.persist(X)
y = client.persist(y)
model_0 = xgb.dask.DaskXGBClassifier(
learning_rate=0.3, random_state=0, n_estimators=4,
tree_method=tree_method)