Update ts_generation_mixin.py

ts_generation_mixin.py

@@ -26,6 +26,7 @@ class TSGenerationMixin(GenerationMixin):
             **model_kwargs,
     ) -> Union[GenerateNonBeamOutput, torch.Tensor]:
         input_ids = input_ids.to(self.device)
+        initial_input_length = input_ids.shape[1]
         if len(input_ids.shape) == 2:
             batch_size, cur_len = input_ids.shape
             if cur_len < self.config.input_token_len:
@@ -37,7 +38,6 @@ class TSGenerationMixin(GenerationMixin):
                 input_ids = input_ids[:, -new_len:]
         else:
             raise ValueError('Input shape must be: [batch_size, seq_len]')
-        initial_input_length = input_ids.shape[1]
 
         # init values
         logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()