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

adapt new fx

Browse Source
This commit is contained in:
oahzxl
2023-01-10 11:56:00 +08:00
parent e532679c95
commit 7ab2db206f
4 changed files with 12 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/autochunk/autochunk_codegen.py
View File

@@ -585,9 +585,9 @@ if CODEGEN_AVAILABLE:
code = "".join(body)
code = "\n".join(" " + line for line in code.split("\n"))
fn_code = f"""
{wrap_stmts}
{wrap_stmts}
{prologue}
{code}"""
{prologue}
{code}"""
# print(fn_code)
return PythonCode(fn_code, globals_)