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

code format

Browse Source
This commit is contained in:
oahzxl
2023-01-06 14:21:49 +08:00
parent 8a634af2f5
commit 2bde9d2b7f
3 changed files with 13 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/autochunk/autochunk_codegen.py
View File

@@ -220,7 +220,9 @@ if CODEGEN_AVAILABLE:
self.max_memory = max_memory
self.meta_node = list(meta_graph.graph.nodes)
# find the chunk regions
self.chunk_region_search = ChunkRegionSearch(meta_graph, max_memory, print_mem)
self.chunk_region_search = ChunkRegionSearch(
meta_graph, max_memory, print_mem
)
self.chunk_infos = self.chunk_region_search.search_region()
def _gen_python_code(