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

casys-kaist-internal · Aug 25, 2023 · a6999e1 · a6999e1
1 parent c1315b9
commit a6999e1
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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