Fix bug that caused ViT models to raise errors.

sony · Mar 4, 2025 · 3254880 · 3254880
1 parent aa2a182
commit 3254880
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/model_compression_toolkit/core/common/graph/memory_graph/compute_graph_max_cut.py b/model_compression_toolkit/core/common/graph/memory_graph/compute_graph_max_cut.py
@@ -27,7 +27,7 @@
 
 def compute_graph_max_cut(memory_graph: MemoryGraph,
                           n_iter: int = 50,
-                          astar_n_iter: int = 500,
+                          astar_n_iter: int = 1000,
                           eps: float = 1e-2) -> Tuple[List[BaseNode], float, List[Cut]]:
     """
     A wrapper function to compute max cut and schedule for a given model.

diff --git a/...on_toolkit/core/pytorch/graph_substitutions/substitutions/scaled_dot_product_attention.py b/...on_toolkit/core/pytorch/graph_substitutions/substitutions/scaled_dot_product_attention.py
@@ -103,7 +103,7 @@ def _get_matmul_node(self, attention_node_name: str, q_node: BaseNode, transpose
         matmul_name = f'{attention_node_name}_matmul1'
         return FunctionalNode(name=matmul_name,
                               framework_attr={},
-                              input_shape=(tuple(q_node.output_shape[0]), tuple(transposed_k_node.output_shape)),
+                              input_shape=(tuple(q_node.output_shape[0]), tuple(transposed_k_node.output_shape[0])),
                               output_shape=tuple(matmul1_output_shape),
                               weights={},
                               layer_class=torch.matmul,