[AutoParallel] fix set_grad_var_shape (#50722)

* fix set_grad_var_shape * recover modify

[AutoParallel] fix set_grad_var_shape (#50722)
* fix set_grad_var_shape * recover modify
76c495d7 · zhaoyingli · GitHub · d4217fc6 · 76c495d7
隐藏空白更改
内联并排

Showing with 2 addition and 0 deletion

python/paddle/distributed/auto_parallel/utils.py python/paddle/distributed/auto_parallel/utils.py +2 -0

未找到文件。
--- a/python/paddle/distributed/auto_parallel/utils.py
+++ b/python/paddle/distributed/auto_parallel/utils.py
@@ -1263,6 +1263,7 @@ def set_grad_var_shape(program, dist_context):
                "relu_grad",
                "exp_grad",
                "sigmoid_grad",
+                "unsqueeze2_grad",
            ]
            forward_list = [
                "reshape2",
@@ -1283,6 +1284,7 @@ def set_grad_var_shape(program, dist_context):
                "relu",
                "exp",
                "sigmoid",
+                "unsqueeze2",
            ]
            if op.type in need_set_shape_list:
                for forward_op in block.ops: