[LLM] fix bug when loss is None in llama modeling.py

PaddlePaddle · May 21, 2024 · 7a53d1b · 7a53d1b
1 parent daf2f3a
commit 7a53d1b
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/paddlenlp/transformers/llama/modeling.py b/paddlenlp/transformers/llama/modeling.py
@@ -1645,8 +1645,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