Update exp_short_term_forecasting.py

salmanali88 · May 17, 2024 · 226cbc5 · 226cbc5
1 parent 57fd251
commit 226cbc5
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/exp/exp_short_term_forecasting.py b/exp/exp_short_term_forecasting.py
@@ -94,6 +94,8 @@ def train(self, setting):
 
                 outputs = self.model(batch_x, None, dec_inp, None)
                 f_dim = -1 if self.args.features == 'MS' else 0
+                outputs = outputs[:, -self.args.pred_len:, f_dim:]
+                batch_y = batch_y[:, -self.args.pred_len:, f_dim:].to(self.device)
 
                 batch_y_mark = batch_y_mark[:, -self.args.pred_len:, f_dim:].to(self.device)
                 loss_value = criterion(batch_x, self.args.frequency_map, outputs, batch_y, batch_y_mark)