From b7c3fc91823e74ff3550e6cb02e7abd474a53dd6 Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Tue, 7 Jun 2022 12:30:41 +0800 Subject: [PATCH] Fix overflow in prediction size. (#7885) (#7980) --- src/c_api/c_api_utils.h | 2 +- tests/python-gpu/test_large_input.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/c_api/c_api_utils.h b/src/c_api/c_api_utils.h index 5a6f5ac5e..bf2913b21 100644 --- a/src/c_api/c_api_utils.h +++ b/src/c_api/c_api_utils.h @@ -115,7 +115,7 @@ inline void CalcPredictShape(bool strict_shape, PredictionType type, size_t rows } } CHECK_EQ( - std::accumulate(shape.cbegin(), shape.cend(), 1, std::multiplies<>{}), + std::accumulate(shape.cbegin(), shape.cend(), static_cast(1), std::multiplies<>{}), chunksize * rows); } diff --git a/tests/python-gpu/test_large_input.py b/tests/python-gpu/test_large_input.py index b99c7a0d9..4c8e06a6f 100644 --- 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)