From 679d8deacc162415e97d991646f50acf4f0a2f9e Mon Sep 17 00:00:00 2001
From: thucpham <minhthuc2502@gmail.com>
Date: Thu, 26 Sep 2024 16:14:57 +0200
Subject: [PATCH] fix crashing on cpu

---
 src/decoding.cc | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/decoding.cc b/src/decoding.cc
index 2884f3bd5..84f39ac37 100644
--- a/src/decoding.cc
+++ b/src/decoding.cc
@@ -526,8 +526,12 @@ namespace ctranslate2 {
 
       disable_tokens.apply();
       std::vector<StorageView> logits_vec;
-      if (return_logits_vocab)
-        logits_vec = build_logits(logits, cur_batch_size * _beam_size);
+      if (return_logits_vocab) {
+        if (is_expanded)
+          logits_vec = build_logits(logits, cur_batch_size * _beam_size);
+        else
+          logits_vec = build_logits(logits, cur_batch_size);
+      }
 
       StorageView log_probs(dtype, device);
       if (bias_towards_prefix) {