mirror of https://github.com/llvm/torch-mlir
[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
parent
bbcdb38d99
commit
2c643adcb9
|
@ -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);
|
||||
|
|
Loading…
Reference in New Issue