Fix a bug for gqa (#101)

* test * test * test
Tencent · Oct 11, 2023 · 1c09bcc · 1c09bcc
1 parent f3d0f2c
commit 1c09bcc
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tencentpretrain/layers/multi_headed_attn.py b/tencentpretrain/layers/multi_headed_attn.py
@@ -94,9 +94,9 @@ def unshape(x):
 
         query, key, value = [linear_layer(x) for linear_layer, x in zip(self.linear_layers, [query, key, value])]
 
-        query = query.view(batch_size, seq_length, heads_num, per_head_size)
-        key = key.view(batch_size, seq_length, self.local_kv_heads_num, per_head_size)
-        value = value.view(batch_size, seq_length, self.local_kv_heads_num, per_head_size)
+        query = query.view(batch_size, -1, heads_num, per_head_size)
+        key = key.view(batch_size, -1, self.local_kv_heads_num, per_head_size)
+        value = value.view(batch_size, -1, self.local_kv_heads_num, per_head_size)
 
         query = query.transpose(1, 2)
         key = repeat_kv(key, self.repeat_num).transpose(1, 2)