From 9fde5666022bd5894406cd01da080eecefe8de0b Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Mon, 4 Nov 2024 22:21:41 +0800 Subject: [PATCH] Fixing linting on main. (#2719) --- server/text_generation_server/models/flash_causal_lm.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/server/text_generation_server/models/flash_causal_lm.py b/server/text_generation_server/models/flash_causal_lm.py index b0085b80..bb908fd0 100644 --- a/server/text_generation_server/models/flash_causal_lm.py +++ b/server/text_generation_server/models/flash_causal_lm.py @@ -1729,9 +1729,11 @@ class FlashCausalLM(Model): # Slots can be discontiguous when prefix caching is enabled, so we need to expand the slot_indices, # then update the slots with the additional indices to ensure we're grabbing the ones that have been # allocated - slot_indices = (batch.slot_indices.unsqueeze(-1).expand(B, new_length) + arange_int).view(-1) + slot_indices = ( + batch.slot_indices.unsqueeze(-1).expand(B, new_length) + arange_int + ).view(-1) slots = batch.slots[slot_indices] - + input_lengths = ( input_lengths.unsqueeze(-1).expand(B, new_length) + arange_int ).view(-1)