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

rename

Browse Source
This commit is contained in:
oahzxl
2023-01-09 17:38:16 +08:00
parent cb9817f75d
commit 1bb1f2ad89
2 changed files with 27 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/autochunk/search_chunk.py
View File

@@ -158,11 +158,11 @@ class SearchChunk(object):
end_trace = output_trace[end_idx]
end_node = self.trace_indice.node_list[end_idx]
chunk_infos = []
for end_dim, _ in enumerate(end_trace["idx"]):
for end_dim, _ in enumerate(end_trace["indice"]):
if len(start_traces) > 1:
continue
for start_node, start_trace in start_traces.items():
for start_dim, _ in enumerate(start_trace["idx"]):
for start_dim, _ in enumerate(start_trace["indice"]):
# dim size cannot be 1
if (
get_node_shape(end_node)[end_dim] == 1