fix: error in BTTF testing case;

WenjieDu · Sep 21, 2023 · e71d1f6 · e71d1f6
1 parent 18945e8
commit e71d1f6
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/tests/forecasting/bttf.py b/tests/forecasting/bttf.py
@@ -38,8 +38,7 @@ class TestBTTF(unittest.TestCase):
     @pytest.mark.xdist_group(name="forecasting-bttf")
     def test_0_forecasting(self):
         predictions = self.bttf.forecast(TEST_SET)
-        logger.info(f"prediction shape: {predictions.shape}")
-        mae = cal_mae(predictions, TEST_SET_INTACT)
+        mae = cal_mae(predictions, TEST_SET_INTACT["X"][:, -N_PRED_STEP:])
         logger.info(f"prediction MAE: {mae}")
 
 

diff --git a/tests/forecasting/config.py b/tests/forecasting/config.py
@@ -18,6 +18,6 @@
 TRAIN_SET = {"X": DATA["train_X"]}
 VAL_SET = {"X": DATA["val_X"]}
 TEST_SET = {"X": DATA["test_X"][:, :-N_PRED_STEP]}
-TEST_SET_INTACT = {"X": DATA["test_X"]}
+TEST_SET_INTACT = {"X": DATA["test_X_intact"]}
 
 RESULT_SAVING_DIR_FOR_CLASSIFICATION = os.path.join(RESULT_SAVING_DIR, "forecasting")