Update modeling_loop_lm.py

modeling_loop_lm.py

@@ -975,7 +975,7 @@ class LoopLMForCausalLM(PreTrainedModel, GenerationMixin):
                 labels.view(-1),
             )
             aux = self.config.lb_loss_factor * lb + self.config.lz_loss_factor * lz
-            loss = ce_loss + aux
+            loss = ce_loss + aux if self.training else ce_loss
 
         return CausalLMOutputWithPast(
             loss=loss,