github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

code style

pull/2364/head
oahzxl 2022-12-23 14:26:43 +08:00
parent d309e9338b
commit 49ba619085
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
chunk_codegen.py
View File

@ -1933,7 +1933,9 @@ def emit_code_with_chunk(
emit_node_func(node, body) emit_node_func(node, body)
# replace input var with chunk var # replace input var with chunk var
for input_node_idx, input_node in enumerate(chunk_inputs[region_idx]): for input_node_idx, input_node in enumerate(chunk_inputs[region_idx]):
for idx, dim in chunk_inputs_dim[region_idx][input_node_idx].items(): for idx, dim in chunk_inputs_dim[region_idx][
input_node_idx
].items():
if idx == node_idx: if idx == node_idx:
chunk_slice = _gen_chunk_slice_dim( chunk_slice = _gen_chunk_slice_dim(
dim, "chunk_idx", _get_node_shape(input_node) dim, "chunk_idx", _get_node_shape(input_node)