fix mypy errors

danielvegamyhre · pytorchmergebot · commit 22c33cace46a · 2025-03-01T16:01:47.000Z
diff --git a/torch/_inductor/fx_passes/micro_pipeline_tp.py b/torch/_inductor/fx_passes/micro_pipeline_tp.py
@@ -445,10 +445,10 @@ def insert_reshape_op(node: torch.fx.Node):
         # Case 2: 3 node match (reshape -> mm -> reshape)
         # - match[0].args[0] will be the "A tensor" input to the reshape op
         # - Has 3D+ shape
-        A_node = match[0].args[0]
-        B_node = mm_node.args[1]
-        A_scale_node = mm_node.args[2]
-        B_scale_node = mm_node.args[3]
+        A_node = cast(torch.fx.Node, match[0].args[0])
+        B_node = cast(torch.fx.Node, mm_node.args[1])
+        A_scale_node = cast(torch.fx.Node, mm_node.args[2])
+        B_scale_node = cast(torch.fx.Node, mm_node.args[3])
 
         A_ndim = _get_tensor(A_node).ndim
         A_scale_ndim = _get_tensor(A_scale_node).ndim
@@ -480,10 +480,10 @@ def insert_reshape_op(node: torch.fx.Node):
 
         return _ScaledMatmul(
             nodes=match,
-            A_node=cast(torch.fx.Node, A_node),
-            B_node=cast(torch.fx.Node, B_node),
-            A_scale_node=cast(torch.fx.Node, A_scale_node),
-            B_scale_node=cast(torch.fx.Node, B_scale_node),
+            A_node=A_node,
+            B_node=B_node,
+            A_scale_node=A_scale_node,
+            B_scale_node=B_scale_node,
             bias_node=get_arg(mm_node, 4, None),
             result_scale_node=get_arg(mm_node, 5, None),
             out_dtype=get_arg(mm_node, 6, None),
