add custmize metrics.

GestaltCogTeam · Oct 1, 2022 · ec6578a · ec6578a
1 parent 407bdd4
commit ec6578a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/step/step_runner/step_runner.py b/step/step_runner/step_runner.py
@@ -1,11 +1,13 @@
 import torch
 
 from basicts.runners import BaseTimeSeriesForecastingRunner
+from basicts.metrics import masked_mae, masked_rmse, masked_mape
 
 
 class STEPRunner(BaseTimeSeriesForecastingRunner):
     def __init__(self, cfg: dict):
         super().__init__(cfg)
+        self.metrics = cfg.get("METRICS", {"MAE": masked_mae, "RMSE": masked_rmse, "MAPE": masked_mape})
         self.forward_features = cfg["MODEL"].get("FROWARD_FEATURES", None)
         self.target_features = cfg["MODEL"].get("TARGET_FEATURES", None)
 

diff --git a/step/step_runner/tsformer_runner.py b/step/step_runner/tsformer_runner.py
@@ -1,4 +1,3 @@
-import time
 import torch
 
 from easytorch.utils.dist import master_only