fix the bug in pseudo_speculative_generate

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

Author: yeyu-nvidia

modelopt/torch/speculative/plugins/megatron_eagle.py

@@ -1441,8 +1441,9 @@ def pseudo_speculative_generate(
 
             # Remove mask tokens from eagle_ids before adding draft_token
             # Remove added hidden_states before
-            eagle_ids = eagle_ids[:, : -self.eagle_config.parallel_draft_step + 1]
-            hidden_states = hidden_states[: -self.eagle_config.parallel_draft_step + 1]
+            if self.eagle_config.parallel_draft_step > 1:
+                eagle_ids = eagle_ids[:, : -self.eagle_config.parallel_draft_step + 1]
+                hidden_states = hidden_states[: -self.eagle_config.parallel_draft_step + 1]
 
             eagle_ids = torch.cat((eagle_ids, draft_token), dim=-1)
             hidden_states = torch.cat(