From e65fcbe58d62c705d2ae26aaf761939124ca21cd Mon Sep 17 00:00:00 2001 From: Saurabh Pathak Date: Mon, 2 Nov 2026 12:32:48 +0530 Subject: [PATCH] fix(sampler): bound top_k to the vocabulary before the int32 cast (#789) --- python/freetoken/engine/sample.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index 708d87642..314ed9440 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -80,7 +80,7 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: # parameters for those rows instead of approximating argmax at low temperature. ts = [1.0 if g else max(p.temperature, MIN_T) for p, g in zip(params, is_greedy)] top_ks = [ - p.top_k if not g and p.top_k >= 1 else self.vocab_size + min(p.top_k, self.vocab_size) if not g and p.top_k >= 1 else self.vocab_size for p, g in zip(params, is_greedy) ] top_ps = [