From e98c8fefd18ecd3a0e4d49905186f5efe5b2d548 Mon Sep 17 00:00:00 2001 From: Xingyou Song Date: Fri, 24 Jan 2025 16:02:05 -0800 Subject: [PATCH] Internal change PiperOrigin-RevId: 719462755 --- optformer/decoding_regression/models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optformer/decoding_regression/models.py b/optformer/decoding_regression/models.py index 108f8ec..1459b10 100644 --- a/optformer/decoding_regression/models.py +++ b/optformer/decoding_regression/models.py @@ -123,7 +123,7 @@ def decode( current_logits[:, ~self._vocab.logit_mask(i)] = NEG_INF # [B, V] - probs = sp.special.softmax(temperature * current_logits, axis=-1) + probs = sp.special.softmax(current_logits / temperature, axis=-1) # Sample tokens. sampled_ids = vectorized_sample(probs)