From b81231c790247620f80a4d2049ed9d57826d872c Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Mon, 4 Nov 2024 15:10:26 +0100 Subject: [PATCH] Fixing linting on main. --- 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)