Typo bug in top p / k sampling (octoml#152)

* Bug fix in top p/k sampling due to typo * remove blank
zxybazh · Jan 10, 2024 · e36aea4 · e36aea4
1 parent 4c56eac
commit e36aea4
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/serve/mlc_serve/model/paged_cache_model.py b/serve/mlc_serve/model/paged_cache_model.py
@@ -178,9 +178,7 @@ def _is_safe_to_sample(prob_like):
         logits_random.div_(t.unsqueeze(dim=1))
 
     if do_top_p or do_top_k:
-        # TODO(vvchernov): looks like there is misprinting. Should logits_random be returned?
-        # If no, where are logits used below?
-        logits = _apply_top_p_top_k(logits_random, top_ps, top_ks)
+        logits_random = _apply_top_p_top_k(logits_random, top_ps, top_ks)
 
     probs = torch.softmax(logits_random, dim=-1)