make predictor tests training on cpu

razdoburdin · Nov 17, 2023 · b62e47e · b62e47e
1 parent 81db75b
commit b62e47e
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/python-sycl/test_sycl_prediction.py b/tests/python-sycl/test_sycl_prediction.py
@@ -125,6 +125,7 @@ def test_sklearn(self):
            tm.make_dataset_strategy(), shap_parameter_strategy)
     @settings(deadline=None)
     def test_shap(self, num_rounds, dataset, param):
+        param.update({"device": "cpu"})
         param = dataset.set_params(param)
         dmat = dataset.get_dmat()
         bst = xgb.train(param, dmat, num_rounds)
@@ -139,11 +140,12 @@ def test_shap(self, num_rounds, dataset, param):
            tm.make_dataset_strategy(), shap_parameter_strategy)
     @settings(deadline=None, max_examples=20)
     def test_shap_interactions(self, num_rounds, dataset, param):
+        param.update({"device": "cpu"})
         param = dataset.set_params(param)
         dmat = dataset.get_dmat()
         bst = xgb.train(param, dmat, num_rounds)
-        bst.set_param({"device": "sycl"})
         test_dmat = xgb.DMatrix(dataset.X, dataset.y, dataset.w, dataset.margin)
+        bst.set_param({"device": "sycl"})
         shap = bst.predict(test_dmat, pred_interactions=True)
         margin = bst.predict(test_dmat, output_margin=True)
         assume(len(dataset.y) > 0)