diff --git a/angelslim/compressor/speculative/train/trainer/eagle3_trainer.py b/angelslim/compressor/speculative/train/trainer/eagle3_trainer.py
index d17c2b98..88b85e25 100644
--- a/angelslim/compressor/speculative/train/trainer/eagle3_trainer.py
+++ b/angelslim/compressor/speculative/train/trainer/eagle3_trainer.py
@@ -208,12 +208,6 @@ def draft_model_training_time_test(
                 target_logits = padding(target_logits, left=False)
                 loss_mask = padding(loss_mask, left=False)
 
-                # Update attention mask to prevent attending to future positions
-                ind = torch.arange(seq_length, device=attention_mask.device)
-                attention_mask[:, :, ind[idx:], ind[: seq_length - idx]] = torch.finfo(
-                    attention_mask.dtype
-                ).min
-
         # Step 8: Compute weighted loss
         ploss_weight = [0.8**i for i in range(len(plosses))]
         ploss = sum([ploss_weight[i] * plosses[i] for i in range(len(plosses))])