llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp
View File

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