sampler: pad dry sequence breakers tensor (#875)

PygmalionAI · Dec 8, 2024 · f0cc35b · f0cc35b
1 parent 9288a98
commit f0cc35b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/aphrodite/modeling/sampling_metadata.py b/aphrodite/modeling/sampling_metadata.py
@@ -791,7 +791,9 @@ def from_lists(cls, temperatures: List[float], dynatemp_mins: List[float],
             pin_memory=pin_memory,
         )
         dry_sequence_breakers_t = torch.tensor(
-            dry_sequence_breaker_ids, 
+            [seq + [0] * (max(len(s) for s in
+                              dry_sequence_breaker_ids) - len(seq)) 
+             for seq in dry_sequence_breaker_ids],
             device="cpu",
             dtype=torch.long,
             pin_memory=pin_memory,