Add "mse" metric option to ALSTM.metric_fn (#1810)

microsoft · Jun 19, 2024 · 73ec0f4 · 73ec0f4
1 parent 155c17f
commit 73ec0f4
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/qlib/contrib/model/pytorch_alstm_ts.py b/qlib/contrib/model/pytorch_alstm_ts.py
@@ -160,6 +160,10 @@ def metric_fn(self, pred, label):
 
         if self.metric in ("", "loss"):
             return -self.loss_fn(pred[mask], label[mask])
+        elif self.metric == "mse":
+            mask = ~torch.isnan(label)
+            weight = torch.ones_like(label)
+            return -self.mse(pred[mask], label[mask], weight[mask])
 
         raise ValueError("unknown metric `%s`" % self.metric)