Update ts_generation_mixin.py

ts_generation_mixin.py

@@ -28,6 +28,8 @@ class TSGenerationMixin(GenerationMixin):
             streamer: Optional["BaseStreamer"] = None,
             **model_kwargs,
     ) -> Union[GenerateNonBeamOutput, torch.Tensor]:
+        input_ids_origin_device = input_ids.device
+        input_ids = input_ids.to(self.device)
         if len(input_ids.shape) == 2:
             batch_size, cur_len = input_ids.shape
         else:
@@ -169,6 +171,7 @@ class TSGenerationMixin(GenerationMixin):
         if streamer is not None:
             streamer.end()
 
+        input_ids.squeeze_(dim=-1).to(input_ids_origin_device)
         if return_dict_in_generate:
             if self.config.is_encoder_decoder:
                 return GenerateEncoderDecoderOutput(
@@ -192,7 +195,7 @@ class TSGenerationMixin(GenerationMixin):
                     past_key_values=model_kwargs.get("past_key_values"),
                 )
         else:
-            return input_ids.squeeze(dim=-1)
+            return input_ids
 
     def _update_model_kwargs_for_generation(
             self,