fix runtime prepare pass (#4502)

Co-authored-by: lufang.chen <lufang.chen@nio.com>
1 year ago · 12c95a9fed
parent 0b00def881
commit 12c95a9fed
1 changed files with 1 additions and 1 deletions
--- a/colossalai/auto_parallel/passes/runtime_preparation_pass.py
+++ b/colossalai/auto_parallel/passes/runtime_preparation_pass.py
@ -144,7 +144,7 @@ def size_value_converting_pass(gm: torch.fx.GraphModule, device_mesh: DeviceMesh
    # DeviceMesh information instructs the scaling of the size value
    device_mesh_info = {}
-    for dim, dim_size in enumerate(device_mesh.mesh_shape):
+    for dim, dim_size in enumerate(device_mesh.shape):
        device_mesh_info[dim] = dim_size
    def _extract_target_dim(node):