fix: fix bug of outputing cum_log_probs in t5 sampling

NVIDIA · Jan 10, 2023 · ab00df0 · ab00df0
1 parent 6ca2d58
commit ab00df0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/fastertransformer/models/t5/T5Decoding.cc b/src/fastertransformer/models/t5/T5Decoding.cc
@@ -921,7 +921,7 @@ void T5Decoding<T>::forward(TensorMap*                 output_tensors,
         }
 
         // Return the cumulative log probability and log probability if requested.
-        if (!using_beam_hyps) {
+        if (beam_width == 1 || !using_beam_hyps) {
             if (output_tensors->isExist("output_log_probs")) {
                 invokeTransposeAxis01(output_tensors->at("output_log_probs").getPtr<float>(),
                                       output_log_probs_buf_,