fix crashing on cpu (#1791)

OpenNMT · Oct 14, 2024 · c687308 · c687308
1 parent 087d085
commit c687308
Showing 1 changed file with 6 additions and 2 deletions.
diff --git 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) {