From ffa66aace080960af3fea338feb578d9b93e25e0 Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Tue, 6 Jul 2021 11:47:17 +0800 Subject: [PATCH] Persist data in dask test. (#7077) --- tests/python/test_with_dask.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py index a83079010..789a66e7f 100644 --- a/tests/python/test_with_dask.py +++ b/tests/python/test_with_dask.py @@ -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)