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

This commit is contained in:
Jiaming Yuan 2022-06-07 12:30:41 +08:00 committed by GitHub
parent 645855e8b1
commit b7c3fc9182
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 4 additions and 2 deletions

View File

@ -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<bst_ulong>(1), std::multiplies<>{}),
chunksize * rows);
}

View File

@ -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)