Update modeling_timer.py

modeling_timer.py

@@ -472,9 +472,9 @@ class TimerForPrediction(TimerPreTrainedModel, TSGenerationMixin):
                     else:
                         output_token_len = h
             lm_head = self.lm_heads[self.output_token_len_map[output_token_len]]
-            predictions = lm_head(hidden_states)
+            predictions = lm_head(hidden_states)[:, -1, :]
             if output_token_len > max_output_length:
-                predictions = predictions[:, :, :max_output_length]
+                predictions = predictions[:, :max_output_length]
             if revin:
                 predictions = predictions * std + mean 
         if not return_dict: