From 4fff77ebcbaadcdd01ad0229a8cb2a56c236db5a Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Tue, 13 Aug 2024 16:25:29 +0200 Subject: [PATCH] Medusa requires reshaping. --- server/text_generation_server/models/flash_causal_lm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/text_generation_server/models/flash_causal_lm.py b/server/text_generation_server/models/flash_causal_lm.py index 13885f28..63c97d94 100644 --- a/server/text_generation_server/models/flash_causal_lm.py +++ b/server/text_generation_server/models/flash_causal_lm.py @@ -1412,7 +1412,7 @@ class FlashCausalLM(Model): ).view(-1) prefix_lens_tensor = ( batch.prefix_lens_tensor.unsqueeze(-1).expand(B, new_length) - ).view(-1) + ).reshape(-1) # Add Copy the block tables for all members block_tables = (