diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py
index 7607af9ac..fec3b00c3 100644
--- a/python/sglang/srt/speculative/eagle_info_v2.py
+++ b/python/sglang/srt/speculative/eagle_info_v2.py
@@ -277,9 +277,8 @@ class EagleVerifyInputV2Mixin:
         device = batch.input_ids.device
 
         candidates = self.draft_token.reshape(bs, self.draft_token_num)
-        predict = torch.zeros(
-            (bs * (self.spec_steps + 1),), dtype=torch.int32, device=device
-        )
+        predict_shape = list(next_token_logits.shape)[:-1]
+        predict = torch.zeros(predict_shape, dtype=torch.int32, device=device).flatten()
         accept_index = torch.full(
             (bs, self.spec_steps + 1), -1, dtype=torch.int32, device=device
         )