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,15 +863,15 @@
     while (lhsSize != rhsSize && currLhsDim < lhsShape.size() &&
            currRhsDim < rhsShape.size()) {
       if (lhsSize < rhsSize) {
-        if (++currLhsDim == lhsShape.size()) {
-          break;
+        currLhsDim++;
+        if (currLhsDim < lhsShape.size()) {
+          lhsSize *= lhsShape[currLhsDim];
         }
-        lhsSize *= lhsShape[currLhsDim];
       } else {
-        if (++currRhsDim == rhsShape.size()) {
-          break;
+        currRhsDim++;
+        if (currRhsDim < rhsShape.size()) {
+          rhsSize *= rhsShape[currRhsDim];
         }
-        rhsSize *= rhsShape[currRhsDim];
       }
     }
     if (lhsSize == rhsSize) {