From 94ca52b7b7801e709cf68aa06de024b361df577b Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Thu, 12 May 2022 02:44:03 +0800 Subject: [PATCH] Fix overflow in prediction size. (#7885) --- 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)