Skip to content

Commit e1f9f80

Browse files
authored
Use gpu predictor for get csr test. (dmlc#8323)
1 parent a71421e commit e1f9f80

File tree

1 file changed

+12
-10
lines changed

1 file changed

+12
-10
lines changed

tests/python-gpu/test_device_quantile_dmatrix.py

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,14 @@
1+
import sys
2+
13
import numpy as np
2-
import xgboost as xgb
34
import pytest
4-
import sys
5-
from hypothesis import given, strategies, settings
6-
from scipy import sparse
5+
from hypothesis import given, settings, strategies
6+
7+
import xgboost as xgb
78

89
sys.path.append("tests/python")
9-
import testing as tm
1010
import test_quantile_dmatrix as tqd
11+
import testing as tm
1112

1213

1314
class TestDeviceQuantileDMatrix:
@@ -107,9 +108,8 @@ def test_ref_dmatrix(self) -> None:
107108
@settings(print_blob=True, deadline=None)
108109
def test_to_csr(self, n_samples, n_features, sparsity) -> None:
109110
import cupy as cp
110-
X, y = tm.make_sparse_regression(
111-
n_samples, n_features, sparsity, False
112-
)
111+
112+
X, y = tm.make_sparse_regression(n_samples, n_features, sparsity, False)
113113
h_X = X.astype(np.float32)
114114

115115
csr = h_X
@@ -130,10 +130,12 @@ def test_to_csr(self, n_samples, n_features, sparsity) -> None:
130130
np.testing.assert_equal(h_ret.indptr, d_ret.indptr)
131131
np.testing.assert_equal(h_ret.indices, d_ret.indices)
132132

133-
booster = xgb.train({"tree_method": "gpu_hist"}, dtrain=d_m)
133+
booster = xgb.train(
134+
{"tree_method": "gpu_hist", "predictor": "gpu_predictor"}, dtrain=d_m
135+
)
134136

135137
np.testing.assert_allclose(
136138
booster.predict(d_m),
137139
booster.predict(xgb.DMatrix(d_m.get_data())),
138-
atol=1e-6
140+
atol=1e-6,
139141
)

0 commit comments

Comments
 (0)