Update ts_generation_mixin.py

ts_generation_mixin.py

@@ -226,12 +226,12 @@ class TSGenerationMixin(GenerationMixin):
             if "decoder_attention_mask" in model_kwargs:
                 decoder_attention_mask = model_kwargs["decoder_attention_mask"]
                 model_kwargs["decoder_attention_mask"] = torch.cat(
-                    [decoder_attention_mask, decoder_attention_mask.new_ones((decoder_attention_mask.shape[0], 1))],
+                    [decoder_attention_mask, decoder_attention_mask.new_ones((decoder_attention_mask.shape[0], horizon_length))],
                     dim=-1,
                 )
 
         if "cache_position" in model_kwargs and model_kwargs["cache_position"] is not None:
-            # model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + horizon_length
-            model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + 1
+            model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + horizon_length
+            # model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + 1
 
         return model_kwargs