Fix overflow in prediction size. (#7885)

2022-05-12 02:44:03 +08:00
parent 8ba4722d04
commit 94ca52b7b7
2 changed files with 4 additions and 2 deletions
--- a/tests/python-gpu/test_large_input.py
+++ b/tests/python-gpu/test_large_input.py
@@ -18,4 +18,6 @@ def test_large_input():
    X = cp.ones((m, n), dtype=np.float32)
    y = cp.ones(m)
    dmat = xgb.DeviceQuantileDMatrix(X, y)
-    xgb.train({"tree_method": "gpu_hist", "max_depth": 1}, dmat, 1)
+    booster = xgb.train({"tree_method": "gpu_hist", "max_depth": 1}, dmat, 1)
+    del y
+    booster.inplace_predict(X)