github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-02 09:38:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

[autoparallel] update getitem handler (#2207)

Browse Source
This commit is contained in:
YuliangLiu0306
2022-12-27 19:58:32 +08:00
committed by GitHub
parent 29868a9ec1
commit 78509124d3
4 changed files with 120 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -223,7 +223,8 @@ def _size_value_converting(gm: torch.fx.GraphModule, device_mesh: DeviceMesh):
node.args = new_args
elif isinstance(getitem_index, (tuple, list)):
assert isinstance(getitem_index[0], slice)
if not isinstance(getitem_index[0], slice):
continue
new_slice_items = []
for slice_item in getitem_index: