Skip to content

GitLab

Unverified Commit 12c95a9f authored by Lufang Chen's avatar Lufang Chen Committed by GitHub
Browse files

fix runtime prepare pass (#4502) 


Co-authored-by: default avatarlufang.chen <lufang.chen@nio.com>
parent 0b00def8
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
colossalai/auto_parallel/passes/runtime_preparation_pass.py
View file @ 12c95a9f
...@@ -144,7 +144,7 @@ def size_value_converting_pass(gm: torch.fx.GraphModule, device_mesh: DeviceMesh ...@@ -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 # DeviceMesh information instructs the scaling of the size value
device_mesh_info = {} 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 device_mesh_info[dim] = dim_size
def _extract_target_dim(node): def _extract_target_dim(node):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment