Attempt #2

huggingface · Oct 23, 2024 · 5c3efbc · 5c3efbc
1 parent 0a01dde
commit 5c3efbc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/server/text_generation_server/models/flash_causal_lm.py b/server/text_generation_server/models/flash_causal_lm.py
@@ -1418,7 +1418,7 @@ def warmup(
                 )
                 max_total_tokens = available_blocks
             else:
-                max_total_tokens = batch.num_blocks
+                max_total_tokens = len(batch.input_ids)
                 max_input_tokens = (
                     batch.num_blocks - 1
                     if max_input_tokens is None