diff --git a/tests/transform_dialect/cuda/mma_elemwise_layout_analysis_codegen_spec.mlir b/tests/transform_dialect/cuda/mma_elemwise_layout_analysis_codegen_spec.mlir
index ecb3410..d1362d0 100644
--- a/tests/transform_dialect/cuda/mma_elemwise_layout_analysis_codegen_spec.mlir
+++ b/tests/transform_dialect/cuda/mma_elemwise_layout_analysis_codegen_spec.mlir
@@ -6,7 +6,8 @@
       // Step 1. Find the fill, matmul and generic ops
       // ===========================================================================
       %fill = transform.structured.match ops{["linalg.fill"]} in %variant_op : (!transform.any_op) -> !transform.any_op
-      %matmul = transform.structured.match ops{["linalg.matmul_transpose_b"]}
+      %matmul = transform.structured.match ops{["linalg.generic"]}
+                  attributes{iterator_types = [#linalg.iterator_type<parallel>, #linalg.iterator_type<parallel>, #linalg.iterator_type<reduction>]}
                   in %variant_op : (!transform.any_op) -> !transform.any_op
       %generic = transform.structured.match ops{["linalg.generic"]}
                   attributes{iterator_types = [#linalg.iterator_type<parallel>, #linalg.iterator_type<parallel>]}
diff --git a/tests/transform_dialect/cuda/mma_reduction_layout_analysis_codegen_spec.mlir b/tests/transform_dialect/cuda/mma_reduction_layout_analysis_codegen_spec.mlir
index d246a9e..a7c32a4 100644
--- a/tests/transform_dialect/cuda/mma_reduction_layout_analysis_codegen_spec.mlir
+++ b/tests/transform_dialect/cuda/mma_reduction_layout_analysis_codegen_spec.mlir
@@ -7,7 +7,8 @@
     // Step 1. Find the fill, matmul and generic ops
     // ===========================================================================
     %fill = transform.structured.match ops{["linalg.fill"]} in %variant_op : (!transform.any_op) -> !transform.any_op
-    %matmul = transform.structured.match ops{["linalg.matmul_transpose_b"]}
+    %matmul = transform.structured.match ops{["linalg.generic"]}
+                attributes{iterator_types = [#linalg.iterator_type<parallel>, #linalg.iterator_type<parallel>, #linalg.iterator_type<reduction>]}
                 in %variant_op : (!transform.any_op) -> !transform.any_op
     %reduce = transform.structured.match ops{["linalg.generic"]}
                 attributes{iterator_types = [#linalg.iterator_type<parallel>, #linalg.iterator_type<reduction>]}