github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-08 12:30:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

[hotfix] update test for latest version (#2060)

Browse Source
This commit is contained in:
YuliangLiu0306
2022-12-02 18:12:30 +08:00
committed by GitHub
parent 19438ea0ef
commit e4293e5077
3 changed files with 15 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
colossalai/auto_parallel/passes/runtime_preparation_pass.py
View File

@@ -126,12 +126,13 @@ def _node_args_converting(gm: torch.fx.GraphModule, device_mesh: DeviceMesh):
if method in (torch.Tensor.view, torch.Tensor.reshape):
for arg in node.args:
if isinstance(arg, Node):
if isinstance(arg._meta_data, int):
if isinstance(arg._meta_data, (int, tuple, list)):
new_args.append(arg._meta_data)
else:
new_args.append(arg)
else:
assert isinstance(arg, int), 'The argument in view node should be either type of Node or int.'
assert isinstance(
arg, (int, tuple, list)), 'The argument in view node should be either type of Node or int.'
new_args.append(arg)
for dim, shard_dims in output_dim_partition_dict.items():