Fix out of range error in SmallVector (CI test before squash)

llvm · Jul 24, 2022 · 2bd287b · 2bd287b
1 parent e6f0529
commit 2bd287b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lib/Conversion/TorchToLinalg/IndirectDataMovement.cpp b/lib/Conversion/TorchToLinalg/IndirectDataMovement.cpp
@@ -386,7 +386,7 @@ class ConvertAtenIndexTensorOp : public OpConversionPattern<AtenIndexTensorOp> {
       resultShape.append(broadcastedIndexShape);
       int j = 0;
       for (auto i : llvm::seq(0, inputRank)) {
-        if (i == indexTensorDims[j]) {
+        if (j < replacedIndexCount && i == indexTensorDims[j]) {
           j++;
           continue;
         }
@@ -456,7 +456,7 @@ class ConvertAtenIndexTensorOp : public OpConversionPattern<AtenIndexTensorOp> {
                   } else {
                     int indexCount = 0, unchanged = 0;
                     for (auto i : llvm::seq(0, inputRank)) {
-                      if (i == indexTensorDims[indexCount]) {
+                      if (indexCount < replacedIndexCount && i == indexTensorDims[indexCount]) {
                         extractionIndices.push_back(
                             castIntToIndex(b, loc, args[indexCount++]));
                         continue;