fix autocast bug

cyber-pioneer · May 9, 2024 · c22bc8b · c22bc8b
1 parent 2619f17
commit c22bc8b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/paddlenlp/generation/utils.py b/paddlenlp/generation/utils.py
@@ -511,7 +511,7 @@ def update_model_kwargs_for_generation(outputs, model_kwargs, is_encoder_decoder
     def update_scores_for_generation(scores, next_scores, length, unfinished_flag):
         # update scores
 
-        unfinished_scores = (scores * length + next_scores) / (length + 1)
+        unfinished_scores = (scores * length.astype(scores.dtype) + next_scores) / (length + 1).astype(scores.dtype)
         scores = paddle.where(unfinished_flag, unfinished_scores, scores)
         return scores