xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-14 10:54:58 +08:00
Code Issues Packages Projects Releases Wiki Activity

Set replacement=True in torch.multinomial (#858)

Browse Source
This commit is contained in:
Woosuk Kwon 2023-08-25 12:22:01 +09:00 committed by GitHub
parent 75c0ca9d43
commit 94d2f59895
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/model_executor/layers/sampler.py
View File

@ -302,7 +302,9 @@ def _sample_from_prompt(
# Random sampling. # Random sampling.
# Sample `best_of` tokens for the prompt. # Sample `best_of` tokens for the prompt.
num_seqs = sampling_params.best_of num_seqs = sampling_params.best_of
next_token_ids = torch.multinomial(prob, num_samples=num_seqs) next_token_ids = torch.multinomial(prob,
num_samples=num_seqs,
replacement=True)
next_token_ids = next_token_ids.tolist() next_token_ids = next_token_ids.tolist()
return next_token_ids return next_token_ids