From 32091a45d7f7b69392ad4f0fc7a212ea4c0bc9e0 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Fri, 25 Aug 2023 02:22:35 +0000 Subject: [PATCH] Add replacement=True --- vllm/model_executor/layers/sampler.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/sampler.py b/vllm/model_executor/layers/sampler.py index 6a50ee59ea67..9717e008972f 100644 --- a/vllm/model_executor/layers/sampler.py +++ b/vllm/model_executor/layers/sampler.py @@ -302,7 +302,9 @@ def _sample_from_prompt( # Random sampling. # Sample `best_of` tokens for the prompt. 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() return next_token_ids