tpu: fix outputs by correcting the next_token_ids shape (#986)

PygmalionAI · Dec 24, 2024 · 0c56d23 · 0c56d23
1 parent 3bb0f07
commit 0c56d23
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/aphrodite/worker/tpu_model_runner.py b/aphrodite/worker/tpu_model_runner.py
@@ -602,7 +602,7 @@ def _execute_model(*args):
                 batch_idx += 1
             else:
                 for seq_id in seq_ids:
-                    next_token_id = next_token_ids[batch_idx][0]
+                    next_token_id = next_token_ids[batch_idx]
                     seq_outputs.append(
                         SequenceOutput(seq_id, next_token_id,
                                        {next_token_id: zero_logprob}))
@@ -723,6 +723,9 @@ def forward(
         sampled_token_ids = torch.multinomial(probs,
                                               num_samples,
                                               replacement=True)
+        if num_samples == 1:
+            argmax_token_ids = argmax_token_ids.squeeze(dim=-1)
+            sampled_token_ids = sampled_token_ids.squeeze(dim=-1)
         next_token_ids = torch.where(t != 0, sampled_token_ids,
                                      argmax_token_ids)
         return next_token_ids