Avoid calling CUDA code on CPU for linear model. (#7154)

2021-09-01 10:45:31 +08:00
parent ba69244a94
commit 3a4f51f39f
4 changed files with 43 additions and 19 deletions
--- a/tests/python-gpu/load_pickle.py
+++ b/tests/python-gpu/load_pickle.py
@@ -19,8 +19,15 @@ class TestLoadPickle:
        assert os.environ['CUDA_VISIBLE_DEVICES'] == '-1'
        bst = load_pickle(model_path)
        x, y = build_dataset()
-        test_x = xgb.DMatrix(x)
-        res = bst.predict(test_x)
+        if isinstance(bst, xgb.Booster):
+            test_x = xgb.DMatrix(x)
+            res = bst.predict(test_x)
+        else:
+            res = bst.predict(x)
+            assert len(res) == 10
+            bst.set_params(n_jobs=1)  # triggers a re-configuration
+            res = bst.predict(x)
+
        assert len(res) == 10

    def test_predictor_type_is_auto(self):