Commit 4a12bc2a authored by Gustaf Ahdritz's avatar Gustaf Ahdritz
Browse files

Fix undefined reference

parent 9617306a
...@@ -439,10 +439,10 @@ class Attention(nn.Module): ...@@ -439,10 +439,10 @@ class Attention(nn.Module):
Returns Returns
[*, Q, C_q] attention update [*, Q, C_q] attention update
""" """
if(use_lma and (q_chunk_size is None or kv_chunk_size is None)): if(use_lma and (lma_q_chunk_size is None or lma_kv_chunk_size is None)):
raise ValueError( raise ValueError(
"If use_lma is specified, q_chunk_size and kv_chunk_size must " "If use_lma is specified, lma_q_chunk_size and "
"be provided" "lma_kv_chunk_size must be provided"
) )
if(use_flash and biases is not None): if(use_flash and biases is not None):
......
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