Skip to content

GitLab

"git@developer.sourcefind.cn:OpenDAS/colossalai.git" did not exist on "f0f35216f1296da701dc0c484e2e1dbe0c05feda"
Commit cb68ee86 authored by oahzxl's avatar oahzxl
Browse files

set benchmark

parent 4d223e18
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
tests/test_autochunk/benchmark_autochunk.py
View file @ cb68ee86
...@@ -98,14 +98,14 @@ def _build_openfold(): ...@@ -98,14 +98,14 @@ def _build_openfold():
def benchmark_evoformer(): def benchmark_evoformer():
# init data and model # init data and model
msa_len = 256 msa_len = 256
pair_len = 256 pair_len = 512
node = torch.randn(1, msa_len, pair_len, 256).cuda() node = torch.randn(1, msa_len, pair_len, 256).cuda()
pair = torch.randn(1, pair_len, pair_len, 128).cuda() pair = torch.randn(1, pair_len, pair_len, 128).cuda()
model = evoformer_base().cuda() model = evoformer_base().cuda()
# build autochunk model # build autochunk model
max_memory = 1000 # MB fit memory mode # max_memory = 1000 # MB, fit memory mode
# max_memory = None # min memory mode max_memory = None # min memory mode
autochunk = _build_autochunk(evoformer_base().cuda(), max_memory, node, pair) autochunk = _build_autochunk(evoformer_base().cuda(), max_memory, node, pair)
# build openfold # build openfold
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment