Update ts_generation_mixin.py

ts_generation_mixin.py

@@ -13,7 +13,7 @@ class TSGenerationMixin(GenerationMixin):
 
     def _greedy_search(
             self,
-            input_ids: torch.LongTensor,
+            input_ids: torch.Tensor,
             logits_processor: Optional[LogitsProcessorList] = None,
             stopping_criteria: Optional[StoppingCriteriaList] = None,
             max_length: Optional[int] = None,
@@ -27,7 +27,11 @@ class TSGenerationMixin(GenerationMixin):
             synced_gpus: bool = False,
             streamer: Optional["BaseStreamer"] = None,
             **model_kwargs,
-    ) -> Union[GenerateNonBeamOutput, torch.LongTensor]:
+    ) -> Union[GenerateNonBeamOutput, torch.Tensor]:
+        if len(input_ids.shape) == 2:
+            batch_size, cur_len = input_ids.shape
+        else:
+            raise ValueError('Input shape must be: [batch_size, seq_len]')
         # init values
         logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
         stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
@@ -82,7 +86,6 @@ class TSGenerationMixin(GenerationMixin):
             )
 
         # keep track of which sequences are already finished
-        batch_size, cur_len = input_ids.shape
         if "inputs_embeds" in model_kwargs:
             cur_len = model_kwargs["inputs_embeds"].shape[1]
         this_peer_finished = False
@@ -189,7 +192,7 @@ class TSGenerationMixin(GenerationMixin):
                     past_key_values=model_kwargs.get("past_key_values"),
                 )
         else:
-            return input_ids
+            return input_ids.squeeze(dim=-1)
 
     def _update_model_kwargs_for_generation(
             self,