Yeyu/set block

modelopt/torch/speculative/eagle/eagle_model.py

@@ -47,5 +47,4 @@ def modify(
         self.eagle_loss_decay_factor = eagle_loss_decay_factor
 
         if eagle_architecture_config.get("parallel_draft_step", 1) > 1:
-            for i in range(eagle_architecture_config.get("parallel_draft_step") - 1):
-                self.register_buffer(f"mask_token_{i}", torch.tensor(-1))
+            self.register_buffer("mask_token", torch.tensor(-1))