debug

Signed-off-by: Ye Yu <[email protected]>

Author: yeyu-nvidia

modelopt/torch/speculative/plugins/megatron_eagle.py

@@ -1284,9 +1284,10 @@ def forward(
 
         # If eagle_freeze_base_model is set to True,
         # the base model is frozen .
-        loss = self.compute_language_model_loss(
-            labels, logits_sbh[:-1] if self.eagle_offline else logits_sbh
-        )
+        if self.eagle_offline:
+            loss = torch.zeros(input_ids.shape).to(input_ids.device)
+        else:
+            loss = self.compute_language_model_loss(labels, logits_sbh)
         loss = 0.0 * loss
 
         if self.eagle_config.parallel_draft_step > 1: