fix: add missing logit index increment in sampling metadata prep (#1059)

aphrodite-engine · Dec 30, 2024 · 0b5588d · 0b5588d
1 parent 525edc1
commit 0b5588d
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/aphrodite/modeling/sampling_metadata.py b/aphrodite/modeling/sampling_metadata.py
@@ -324,6 +324,7 @@ def sample(logits):
             sample_indices.extend(range(logit_idx, logit_idx + sample_len))
             categorized_sample_indices[sampling_params.sampling_type].extend(
                 list(range(logit_idx, logit_idx + sample_len)))
+            logit_idx += sample_len
 
         if cache is not None:
             sample_obj.sampling_params = sampling_params