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

[NFC] polish colossalai/pipeline/utils.py code style (#1562)

pull/1550/head
Zirui Zhu 2022-09-08 16:25:57 +08:00 committed by Frank Lee
parent e070ca45c6
commit f566c9b98d
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/pipeline/utils.py
View File

@ -250,6 +250,7 @@ def call_module(module, args=None, kwargs=None):
else:
return module(*args_needed, **kwargs)
def customized_partition(exec_seq):
'''
This function will analyze the exec_seq. In the exec_seq, users will use 'SPLIT_NODE' as an
@ -269,4 +270,3 @@ def customized_partition(exec_seq):
stop += 1
customized_parts[rank] = [(start, stop)]
return customized_parts