fix msc testcase

apache · Sep 17, 2024 · 43268e1 · 43268e1
1 parent f6f6c1a
commit 43268e1
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/tests/python/contrib/test_msc/test_graph_build.py b/tests/python/contrib/test_msc/test_graph_build.py
@@ -2363,13 +2363,13 @@ def forward(self, q_data, k_data, v_data):
         ],
         "outputs": [
             {
-                "name": "attention",
-                "shape": [1, seq, 8, 64],
+                "name": "permute_dims_3",
+                "shape": [1, 8, seq, 64],
                 "dtype": "float32",
-                "layout": "ABCD",
+                "layout": "ACBD",
             }
         ],
-        "nodes": {"total": 4, "input": 3, "msc.attention": 1},
+        "nodes": {"total": 5, "input": 3, "msc.attention": 1, "permute_dims": 1},
     }
     if dynamic:
         expected1["prims"] = {"total": 1, "shape": 1}
@@ -2395,13 +2395,13 @@ def forward(self, q_data, k_data, v_data, mask):
         ],
         "outputs": [
             {
-                "name": "attention_bias",
-                "shape": [1, seq, 8, 64],
+                "name": "permute_dims_3",
+                "shape": [1, 8, seq, 64],
                 "dtype": "float32",
-                "layout": "ABCD",
+                "layout": "ACBD",
             }
         ],
-        "nodes": {"total": 5, "input": 4, "msc.attention": 1},
+        "nodes": {"total": 6, "input": 4, "msc.attention": 1, "permute_dims": 1},
     }
     if dynamic:
         expected2["prims"] = {"total": 1, "shape": 1}