Skip to content

Commit c687308

Browse files
authored
fix crashing on cpu (#1791)
1 parent 087d085 commit c687308

File tree

1 file changed

+6
-2
lines changed

1 file changed

+6
-2
lines changed

src/decoding.cc

+6-2
Original file line numberDiff line numberDiff line change
@@ -526,8 +526,12 @@ namespace ctranslate2 {
526526

527527
disable_tokens.apply();
528528
std::vector<StorageView> logits_vec;
529-
if (return_logits_vocab)
530-
logits_vec = build_logits(logits, cur_batch_size * _beam_size);
529+
if (return_logits_vocab) {
530+
if (is_expanded)
531+
logits_vec = build_logits(logits, cur_batch_size * _beam_size);
532+
else
533+
logits_vec = build_logits(logits, cur_batch_size);
534+
}
531535

532536
StorageView log_probs(dtype, device);
533537
if (bias_towards_prefix) {

0 commit comments

Comments
 (0)