diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py
index 2b3e7abad..b3c54736e 100644
--- a/vllm_omni/entrypoints/async_omni.py
+++ b/vllm_omni/entrypoints/async_omni.py
@@ -360,7 +360,6 @@ class AsyncOmni(OmniBase):
                         raise RuntimeError(result)  # Request Finished due to error
 
                     engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
-                    stage.set_engine_outputs(engine_outputs)
                     if isinstance(engine_outputs, list):
                         engine_outputs = engine_outputs[0]
                     finished = engine_outputs.finished
@@ -419,7 +418,9 @@ class AsyncOmni(OmniBase):
                                 final_output_type=stage.final_output_type,
                                 request_output=engine_outputs,
                             )
-
+                if not isinstance(engine_outputs, list):
+                    engine_outputs = [engine_outputs]
+                stage.set_engine_outputs(engine_outputs)
                 # Forward to next stage if there is one
                 next_stage_id = stage_id + 1
                 if next_stage_id <= final_stage_id_for_e2e and finished:
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
index f2da62b4a..05169303a 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
@@ -651,7 +651,11 @@ class Qwen3OmniMoeForConditionalGeneration(
             )
             inputs_embeds = summed_embeddings.clone()
         else:
-            code_predictor_codes = torch.zeros((0, self.talker.num_code_groups), dtype=torch.long)
+            code_predictor_codes = torch.zeros(
+                (inputs_embeds.shape[0], self.talker.num_code_groups),
+                device=self._module_device(self.talker),
+                dtype=torch.long,
+            )
         inputs_embeds = (inputs_embeds + text_step).reshape(-1, self.talker_config.text_config.hidden_size)
         return inputs_embeds, code_predictor_codes.squeeze(-1).detach().to("cpu").contiguous()