Fixing linting on main.

huggingface · Nov 4, 2024 · b81231c · b81231c
1 parent aadc9cb
commit b81231c
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/server/text_generation_server/models/flash_causal_lm.py b/server/text_generation_server/models/flash_causal_lm.py
@@ -1729,9 +1729,11 @@ def forward(
             # 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)