Fix binary search in bisect_percentile_op (#73146)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/73146

Binary search can overflow; this fixes it.

Test Plan: Sandcastle

Reviewed By: meyering

Differential Revision: D34365186

fbshipit-source-id: f92a810b49ef5ce345d0b019b584fe3c1f5ae017
(cherry picked from commit 9c2133ec6f)
This commit is contained in:
Richard Barnes 2022-02-21 13:49:19 -08:00 committed by PyTorch MergeBot
parent 5dad19fef0
commit c2255c36ec

View File

@ -115,13 +115,10 @@ class BisectPercentileOp final : public Operator<Context> {
int lo,
int hi,
float val) {
int mid;
bool low_cond, high_cond;
while (lo < hi) {
mid = (lo + hi) >> 1;
low_cond = (data[mid] <= val);
high_cond = (val < data[mid + 1]);
const auto mid = lo + (hi - lo) / 2;
const bool low_cond = (data[mid] <= val);
const bool high_cond = (val < data[mid + 1]);
if (low_cond && high_cond) {
return mid;
} else if (!low_cond) {