fix(server): fix cohere (#2249)

2024-07-18 14:00:13 +00:00 · 2024-07-18 14:00:13 +00:00 · 1d1b1efa01
parent da82c63a4f
commit 1d1b1efa01
1 changed files with 1 additions and 1 deletions
--- a/server/text_generation_server/models/custom_modeling/flash_cohere_modeling.py
+++ b/server/text_generation_server/models/custom_modeling/flash_cohere_modeling.py
@ -259,8 +259,8 @@ class FlashCohereAttention(torch.nn.Module):
        cu_seqlen_prefill,
        kv_cache,
        block_tables,
-        input_lengths,
        slots,
+        input_lengths,
        max_s,
    ):
        qkv = self.query_key_value(hidden_states)