From e1d30a0dbd65b1706523cf9b38843d075f5c4928 Mon Sep 17 00:00:00 2001 From: lilujia Date: Fri, 17 May 2024 12:38:55 +0800 Subject: [PATCH] [LLM] fix bug when loss is None in llama modeling.py --- paddlenlp/transformers/llama/modeling.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/paddlenlp/transformers/llama/modeling.py b/paddlenlp/transformers/llama/modeling.py index 8f2dd1c36415..78fc26d08dbf 100755 --- a/paddlenlp/transformers/llama/modeling.py +++ b/paddlenlp/transformers/llama/modeling.py @@ -1654,8 +1654,11 @@ def forward(self, prediction_scores, masked_lm_labels): binary_sequence = paddle.where( masked_lm_loss > 0, paddle.ones_like(masked_lm_loss), paddle.zeros_like(masked_lm_loss) ) - sum_ = paddle.sum(binary_sequence) - loss = 0 if sum_ == 0 else paddle.sum(masked_lm_loss * binary_sequence) / sum_ + count = paddle.sum(binary_sequence) + if count == 0: + loss = paddle.sum(masked_lm_loss * binary_sequence) + else: + loss = paddle.sum(masked_lm_loss * binary_sequence) / count return loss