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

PaddlePaddle · May 17, 2024 · 2d2a503 · 2d2a503
1 parent daf2f3a
commit 2d2a503
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/paddlenlp/transformers/llama/modeling.py b/paddlenlp/transformers/llama/modeling.py
@@ -1646,7 +1646,11 @@ def forward(self, prediction_scores, masked_lm_labels):
                 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_
+            if sum_ == 0:
+                loss = paddle.zeros([], dtype=masked_lm_loss.dtype)
+                loss.stop_gradient = False
+            else:
+                loss = paddle.sum(masked_lm_loss * binary_sequence) / sum_
 
         return loss