def kv_seq_len

huggingface · Oct 9, 2023 · ae73232 · ae73232
1 parent 34b7ab1
commit ae73232
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/optimum/bettertransformer/models/attention.py b/optimum/bettertransformer/models/attention.py
@@ -695,6 +695,7 @@ def gpt_bigcode_wrapped_scaled_dot_product(
     # MHA models: (batch_size, num_heads, query_length, head_dim)
     query_shape = query.shape
     batch_size = query_shape[0]
+    kv_seq_len = key.shape[-2]
 
     if self.multi_query:
         query_length = query_shape[1]