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

seperate trace flow

Browse Source
This commit is contained in:
oahzxl
2023-01-06 17:24:23 +08:00
parent 4748967fb1
commit a6cdbf9161
6 changed files with 447 additions and 424 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/autochunk/autochunk_codegen.py
View File

@@ -167,7 +167,7 @@ def emit_code_with_chunk(
)
# ones like
if "ones_like" in node.name:
meta_node = chunk_region_search.index_tracer.node_list[node_idx]
meta_node = chunk_region_search.trace_index.node_list[node_idx]
chunk_dim = chunk_infos[region_idx]["node_chunk_dim"][meta_node][
"chunk_dim"
]