add grad var shape infer of matmul and mul (#45237)

PaddlePaddle · Aug 18, 2022 · 520b054 · 520b054
1 parent 41294cb
commit 520b054
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/paddle/distributed/auto_parallel/utils.py b/python/paddle/distributed/auto_parallel/utils.py
@@ -1043,7 +1043,7 @@ def set_grad_var_shape(program, dist_context):
                     "fill_zeros_like"
             ]:
                 forward_var_name = op.input_arg_names[0]
-            elif op.type == "matmul_v2_grad":
+            elif op.type == "matmul_v2_grad" or op.type == "matmul_grad" or op.type == "mul_grad":
                 forward_var_name = None
                 for output_name in op.output_names:
                     if var_name in op.output(output_name):