Commit 805b45cc authored by Gustaf Ahdritz's avatar Gustaf Ahdritz
Browse files

Fix bug in template stack

parent 3f57b4a0
...@@ -393,7 +393,7 @@ class TemplatePairStack(nn.Module): ...@@ -393,7 +393,7 @@ class TemplatePairStack(nn.Module):
) )
blocks = [ blocks = [
partial(b, partial(b,
chunk_size=chunk_size, chunk_size=tuned_chunk_size,
_attn_chunk_size=max(chunk_size, tuned_chunk_size // 4), _attn_chunk_size=max(chunk_size, tuned_chunk_size // 4),
) for b in blocks ) for b in blocks
] ]
......
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