fix(server): fix cohere (#2249)

1d1b1efa · OlivierDehaene · GitHub · da82c63a · 1d1b1efa
Unverified Commit 1d1b1efa authored Jul 18, 2024 by OlivierDehaene Committed by GitHub Jul 18, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

server/text_generation_server/models/custom_modeling/flash_cohere_modeling.py ...on_server/models/custom_modeling/flash_cohere_modeling.py +1 -1

No files found.
--- 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)