fix ut

Signed-off-by: Yu, Zhentao <[email protected]>
intel · Mar 25, 2024 · a8f8cc7 · a8f8cc7
1 parent 804a556
commit a8f8cc7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/neural_speed/__init__.py b/neural_speed/__init__.py
@@ -393,7 +393,7 @@ def __call__(self, model_input, reinit=False, logits_all=False, **kwargs):
                         raw_logits[i] = np.vstack((padding_row * (-np.inf), raw_logits[i]))
                     else:
                         raw_logits[i] = np.vstack((raw_logits[i], padding_row * (-np.inf)))
-            return raw_logits
+            return np.array(raw_logits)
         else:
             print("Please input torch.Tensor")
         return

diff --git a/neural_speed/models/model_utils/scheduler.cpp b/neural_speed/models/model_utils/scheduler.cpp
@@ -189,7 +189,7 @@ bool Cont_batch_gen_worker::update_seqs(std::vector<sequence>* seqs, const int&
     }
     if (!m_ctx->beam_search) {
       if (next_tokens.size() != n_input) {
-        fprintf(stderr, "%s: error: wrong next_tokens size %d, which should be %d.\n", __func__, next_tokens.size(),
+        fprintf(stderr, "%s: error: wrong next_tokens size %ld, which should be %d.\n", __func__, next_tokens.size(),
                 n_input);
         return false;
       }