repair dist_scale (#50674)

c1b5e7c2 · xu98bin · GitHub · 2f4763ee · c1b5e7c2
隐藏空白更改
内联并排

Showing with 2 addition and 0 deletion

python/paddle/distributed/auto_parallel/operators/dist_scale.py .../paddle/distributed/auto_parallel/operators/dist_scale.py +2 -0

未找到文件。
--- a/python/paddle/distributed/auto_parallel/operators/dist_scale.py
+++ b/python/paddle/distributed/auto_parallel/operators/dist_scale.py
@@ -74,6 +74,8 @@ class DistributedScaleImpl(DistributedOperatorImpl):
                [x_dims_mapping, out_dims_mapping], [i, i]
            )
            if dim_changed:
+                op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
+                op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
                changed = True

        if changed: