Fix bug that caused ViT models to raise errors.

Author: liord

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.

model_compression_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,