Fix overflow in prediction size. (#7885) (#7980)

2022-06-07 12:30:41 +08:00
parent 645855e8b1
commit b7c3fc9182
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)