diff --git a/deepspeed/runtime/hybrid_engine.py b/deepspeed/runtime/hybrid_engine.py
index cff4219d916b..da6f7a9be54e 100644
--- a/deepspeed/runtime/hybrid_engine.py
+++ b/deepspeed/runtime/hybrid_engine.py
@@ -223,7 +223,7 @@ def generate(self, *inputs, **kwargs):
                 dist.all_gather_into_tensor(output, input_cont, group=self.mp_group)
 
                 if len(inputs) > 0:
-                    inputs = (output, )
+                    inputs = (output, *inputs[1:])
                 else:
                     kwargs['input_ids'] = output