Fix typo in min_p sampling (#385)

state-spaces · Jun 11, 2024 · 8f42a5e · 8f42a5e
1 parent 219f03c
commit 8f42a5e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mamba_ssm/utils/generation.py b/mamba_ssm/utils/generation.py
@@ -105,7 +105,7 @@ def sample(logits, top_k=1, top_p=0.0, min_p=0.0, temperature=1.0):
                 logits_top = logits.clone()
                 max_prob = logits_top[..., 0].item()
                 min_prob = max_prob * min_p
-                modify_logits_for_min_p_filtering(logits_top, min_p)
+                modify_logits_for_min_p_filtering(logits_top, min_prob)
                 if temperature != 1.0:
                     logits_top /= temperature
                 return torch.multinomial(torch.softmax(logits_top, dim=-1), num_samples=1).squeeze(dim=-1)