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 17:42:51 +08:00
parent 9b1b890347
commit 786a398a6b
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
chunk_codegen.py
View File

@ -920,9 +920,13 @@ class IndexTracer(object):
# loop cur node's all arg until out of chunk # loop cur node's all arg until out of chunk
while len(tmp_cur_prepose_nodes) > 0: while len(tmp_cur_prepose_nodes) > 0:
if prepose_flag == False:
break
tmp_next_prepose_nodes = [] tmp_next_prepose_nodes = []
tmp_cur_related_prepose_nodes.extend(tmp_cur_prepose_nodes) tmp_cur_related_prepose_nodes.extend(tmp_cur_prepose_nodes)
for cur_prepose_node in tmp_cur_prepose_nodes: for cur_prepose_node in tmp_cur_prepose_nodes:
if prepose_flag == False:
break
for cur_prepose_node_arg in cur_prepose_node.args: for cur_prepose_node_arg in cur_prepose_node.args:
if type(cur_prepose_node_arg) != type(cur_prepose_node): if type(cur_prepose_node_arg) != type(cur_prepose_node):
continue continue
@ -942,8 +946,6 @@ class IndexTracer(object):
else: else:
prepose_flag = False prepose_flag = False
break break
break
break
# non compute op # non compute op
else: else:
tmp_next_prepose_nodes.append(cur_prepose_node_arg) tmp_next_prepose_nodes.append(cur_prepose_node_arg)