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

rename trace_index to trace_indice

Browse Source
This commit is contained in:
oahzxl
2023-01-09 17:25:13 +08:00
parent 065f0b4c27
commit 0ea903b94e
6 changed files with 74 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/autochunk/autochunk_codegen.py
View File

@@ -94,9 +94,9 @@ def _replace_reshape_size(context, node_name, reshape_size_dict):
return context
def _replace_ones_like(search_chunk, chunk_infos, region_idx, node_idx, node, body):
def _replace_ones_like(search_chunk: SearchChunk, chunk_infos, region_idx, node_idx, node, body):
if "ones_like" in node.name:
meta_node = search_chunk.trace_index.node_list[node_idx]
meta_node = search_chunk.trace_indice.node_list[node_idx]
chunk_dim = chunk_infos[region_idx]["node_chunk_dim"][meta_node]["chunk_dim"]
if get_node_shape(meta_node)[chunk_dim] != 1:
source_node = meta_node.args[0].args[0]