Update vllm/v1/sample/sampler.py

Co-authored-by: Nick Hill <[email protected]>
vllm-project · Jan 22, 2025 · a49ab7e · a49ab7e
1 parent 832d5c8
commit a49ab7e
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py
@@ -140,9 +140,8 @@ def get_logprobs(
         topk_indices = topk_indices.to(torch.int32)
 
         # Concatenate with the token_ids
-        sampled_logprobs = logprobs[torch.arange(logprobs.size(0)),
-                                    token_ids].unsqueeze(-1)
         token_ids = token_ids.unsqueeze(-1)
+        sampled_logprobs = logprobs.gather(-1, token_ids)
         topk_indices = torch.cat([token_ids, topk_indices], dim=1)
         topk_logprobs = torch.cat([sampled_logprobs, topk_logprobs], dim=1)