[TORCH][DECOMPOSE] Fix bug in computeReductionType API

-- This commit fixes a bug in computeReductionType API.
-- The bug pertains to removal of `dim` from the `sizes` array.

Signed-off-by: Abhishek Varma <abhishek@nod-labs.com>
pull/1633/head
Abhishek Varma 2022-11-23 09:36:44 +00:00 committed by Vivek Khandelwal
parent bbcdb38d99
commit 2c643adcb9
1 changed files with 1 additions and 1 deletions

View File

@ -59,7 +59,7 @@ static Type computeReductionType(PatternRewriter &rewriter, Operation *op,
if (keepDim)
sizes[dimInt] = 1;
else
sizes.erase(sizes.begin() + dimInt - 1);
sizes.erase(sizes.begin() + dimInt);
} else {
unsigned reducedRank = keepDim ? inputRank : inputRank - 1;
sizes.resize(reducedRank, kUnknownSize);