From a8f91854f2085126b54172c9b99f48067abb1afe Mon Sep 17 00:00:00 2001 From: "lufang.chen" Date: Wed, 23 Aug 2023 07:25:06 +0000 Subject: [PATCH] fix runtime prepare pass --- colossalai/auto_parallel/passes/runtime_preparation_pass.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/colossalai/auto_parallel/passes/runtime_preparation_pass.py b/colossalai/auto_parallel/passes/runtime_preparation_pass.py index 1a6dc7815176..0ed0742ee57e 100644 --- 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):