diff --git a/internlm/core/scheduler/pipeline_scheduler.py b/internlm/core/scheduler/pipeline_scheduler.py
index 055bb92..a9955a2 100644
--- a/internlm/core/scheduler/pipeline_scheduler.py
+++ b/internlm/core/scheduler/pipeline_scheduler.py
@@ -1342,7 +1342,7 @@ class InterleavedPipelineScheduler(PipelineScheduler):
         accum_moe_loss = self._accum_moe_loss
 
         accum_loss = self._accum_loss
-        if return_loss:
+        if accum_loss is not None:
             accum_loss += self._accum_moe_loss
 
         self._clear_state()