diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -863,10 +863,14 @@
     while (lhsSize != rhsSize && currLhsDim < lhsShape.size() &&
            currRhsDim < rhsShape.size()) {
       if (lhsSize < rhsSize) {
-        currLhsDim++;
+        if (++currLhsDim == lhsShape.size()) {
+          break;
+        }
         lhsSize *= lhsShape[currLhsDim];
       } else {
-        currRhsDim++;
+        if (++currRhsDim == rhsShape.size()) {
+          break;
+        }
         rhsSize *= rhsShape[currRhsDim];
       }
     }