Update modeling_qwen_yarn.py

modeling_qwen_yarn.py

@@ -1156,6 +1156,10 @@ class QWenLMHeadModel(QWenPreTrainedModel):
             output = (lm_logits,) + transformer_outputs[1:]
             return ((loss,) + output) if loss is not None else output
 
+        #训练时节约显存
+        # if self.training:
+            # lm_logits=None
+            
         return CausalLMOutputWithPast(
             loss=loss,
             logits=lm_logits,