Set device in device dmatrix. (#5596)

2020-04-25 13:42:53 +08:00
parent ef26bc45bf
commit e726dd9902
8 changed files with 41 additions and 5 deletions
--- a/tests/python-gpu/test_gpu_prediction.py
+++ b/tests/python-gpu/test_gpu_prediction.py
@@ -121,6 +121,7 @@ class TestGPUPredict(unittest.TestCase):
    @pytest.mark.skipif(**tm.no_cupy())
    def test_inplace_predict_cupy(self):
        import cupy as cp
+        cp.cuda.runtime.setDevice(0)
        rows = 1000
        cols = 10
        cp_rng = cp.random.RandomState(1994)