Merge pull request #20 from dansuissa/fix-logits-rescaling

Fix multiple temperature scaling in sample_non_deterministic()
Adibvafa · Feb 6, 2025 · aaa68b4 · aaa68b4
2 parents 2842ef0 + c5276a6
commit aaa68b4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/CodonTransformer/CodonPrediction.py b/CodonTransformer/CodonPrediction.py
@@ -284,8 +284,8 @@ def sample_non_deterministic(
         raise ValueError("top_p must be a float between 0 and 1.")
 
     # Compute probabilities using temperature scaling
-    logits /= temperature
-    probs = torch.softmax(logits, dim=-1)
+    probs = torch.softmax(logits / temperature, dim=-1)
+
 
     # Remove batch dimension if present
     if probs.dim() == 3: