From 79780acad0a23e17a954d9764a69b5789ed19202 Mon Sep 17 00:00:00 2001
From: Valery Chernov <valery.chernov@deelvin.com>
Date: Fri, 2 Feb 2024 20:23:27 +0400
Subject: [PATCH] fixes

---
 serve/mlc_serve/model/model_common.py | 2 +-
 serve/mlc_serve/model/tvm_model.py    | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/serve/mlc_serve/model/model_common.py b/serve/mlc_serve/model/model_common.py
index 3c9add6618..feb3b61a99 100644
--- a/serve/mlc_serve/model/model_common.py
+++ b/serve/mlc_serve/model/model_common.py
@@ -455,7 +455,7 @@ def sample_loglikelihood_from_logits(
                     sequence_id=sequence_id,
                     generated_tokens=[],
                     error=None,
-                    logprob_info=logprob_infos[i],
+                    logprob_info=get_logprob_infos(i, logprob_infos),
                 )
             )
 
diff --git a/serve/mlc_serve/model/tvm_model.py b/serve/mlc_serve/model/tvm_model.py
index e161f861d3..b825aac53f 100644
--- a/serve/mlc_serve/model/tvm_model.py
+++ b/serve/mlc_serve/model/tvm_model.py
@@ -348,7 +348,7 @@ def generate(
                     input_ids,
                     positions,
                     seq_lens,
-                    cache.cache_blocks,
+                    self.cache_blocks,
                     slot_mapping,
                     all_logits,
                     self.params,