github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-17 23:18:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

seperate trace flow

Browse Source
This commit is contained in:
oahzxl
2023-01-06 17:24:23 +08:00
parent 4748967fb1
commit a6cdbf9161
6 changed files with 447 additions and 424 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_autochunk/benchmark_autochunk.py
View File

@@ -104,8 +104,8 @@ def benchmark_evoformer():
model = evoformer_base().cuda()
# build autochunk model
# max_memory = 1000 # MB fit memory mode
max_memory = None # min memory mode
max_memory = 1000 # MB fit memory mode
# max_memory = None # min memory mode
autochunk = _build_autochunk(evoformer_base().cuda(), max_memory, node, pair)
# build openfold