huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 12:24:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

add contiguous

Browse Source
This commit is contained in:
OlivierDehaene 2024-04-04 18:48:58 +02:00
parent 5088005908
commit 4a02d3505f
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
server/text_generation_server/models/custom_modeling/flash_cohere_modeling.py
View File

@ -222,8 +222,8 @@ class FlashCohereAttention(torch.nn.Module):
dim=1,
)
if self.use_qk_norm:
query = self.q_norm(query)
key = self.k_norm(key)
query = self.q_norm(query.contiguous())
key = self.k_norm(key.contiguous())
query = query.view(-1, self.num_heads, self.head_size)
key = key.view(-1, self.num_key_value_heads, self.head_size)