Update SpecTree.py

Infini-AI-Lab · Mar 22, 2024 · 12844ff · 12844ff
1 parent 61bd16e
commit 12844ff
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/Tree/SpecTree.py b/Tree/SpecTree.py
@@ -149,7 +149,7 @@ def accept_step(self, parent_id :int):
             q = softmax(draft_logits / self.temperature, dim=-1)
             r = self.r[pos + (self.ground_truth_len - 1)]
 
-            if p[token] > r * q[token]:
+            if p[token] >= r * q[token]:
                 return (pos + (self.ground_truth_len - 1), None)
             else:
                 p = self.residual_graph(p, q)
@@ -486,4 +486,4 @@ def verbose(self):
 
 
 
-
+