Internal change

PiperOrigin-RevId: 719462755
google-research · Jan 31, 2025 · e98c8fe · e98c8fe
1 parent 6c9409a
commit e98c8fe
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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)