Fix unsqueeze optimize pass

Differential Revision: D69812661 Pull Request resolved: #8564
pytorch · Feb 19, 2025 · e5dc18a · e5dc18a
1 parent b6ffe1a
commit e5dc18a
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/backends/transforms/view_copy_to_squeeze_unsqueeze.py b/backends/transforms/view_copy_to_squeeze_unsqueeze.py
@@ -75,7 +75,11 @@ def find_unsqueeze_dim(
         j = 0
         idx = -1
         while j < len(view_shape):
-            if input_shape[i] != view_shape[j]:
+            # account for added dim being last dim in view_shape
+            if i == j and j == len(input_shape):
+                if view_shape[j] != 1:
+                    return None
+            elif input_shape[i] != view_shape[j]:
                 if view_shape[j] == 1:
                     idx = j
                     i -= 1