diff --git a/miles/rollout/inference_rollout/inference_rollout_common.py b/miles/rollout/inference_rollout/inference_rollout_common.py
index ce1ffee73f7..2fed07baab2 100644
--- a/miles/rollout/inference_rollout/inference_rollout_common.py
+++ b/miles/rollout/inference_rollout/inference_rollout_common.py
@@ -126,7 +126,7 @@ async def generate_and_rm_group(
     sampling_params: dict[str, Any],
     evaluation: bool = False,
     sample_done_callback: Callable[[], None] | None = None,
-) -> list[Sample]:
+) -> list[Sample | list[Sample]]:
     args = state.args
 
     if state.aborted:
diff --git a/miles/rollout/inference_rollout/inference_rollout_train.py b/miles/rollout/inference_rollout/inference_rollout_train.py
index a486bb2fbc3..5f390a267a5 100644
--- a/miles/rollout/inference_rollout/inference_rollout_train.py
+++ b/miles/rollout/inference_rollout/inference_rollout_train.py
@@ -119,7 +119,7 @@ async def generate_rollout_async(
         logger.debug(f"[rollout] asyncio.wait returned: {len(done)} done, {len(pendings)} pending")
         for task in done:
             try:
-                group: list[Sample] = task.result()
+                group: list[Sample | list[Sample]] = task.result()
             except Exception as e:
                 logger.error(f"[rollout] Task raised exception: {e!r}", exc_info=True)
                 continue