fix: run pre commit lints

This commit is contained in:
drbh 2024-11-01 12:11:57 -04:00
parent 3bafa0eb7b
commit a604bfe450
1 changed files with 4 additions and 2 deletions

View File

@ -1723,7 +1723,9 @@ 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 = (