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_from_cupy.py
+++ b/tests/python-gpu/test_from_cupy.py
@@ -136,3 +136,14 @@ Arrow specification.'''
        n = 100
        X = cp.random.random((n, 2))
        xgb.DeviceQuantileDMatrix(X.toDlpack())
+
+    @pytest.mark.skipif(**tm.no_cupy())
+    @pytest.mark.mgpu
+    def test_specified_device(self):
+        import cupy as cp
+        cp.cuda.runtime.setDevice(0)
+        dtrain = dmatrix_from_cupy(
+            np.float32, xgb.DeviceQuantileDMatrix, np.nan)
+        with pytest.raises(xgb.core.XGBoostError):
+            xgb.train({'tree_method': 'gpu_hist', 'gpu_id': 1},
+                      dtrain, num_boost_round=10)