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

close mem and code print

Browse Source
This commit is contained in:
oahzxl
2023-01-06 14:19:45 +08:00
parent 1a6d2a740b
commit 8a634af2f5
3 changed files with 10 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_autochunk/benchmark_autochunk.py
View File

@@ -64,7 +64,7 @@ def _build_autochunk(model, max_memory, node, pair):
)
# set code_gen
codegen = AutoChunkCodeGen(gm_prop, max_memory)
codegen = AutoChunkCodeGen(gm_prop, max_memory, print_mem=False)
graph.set_codegen(codegen)
gm = ColoGraphModule(model, graph)
gm.recompile()