'fix-bug'

GestaltCogTeam · Jan 2, 2025 · f4bd9f3 · f4bd9f3
2 parents f4ef802 + 489884e
commit f4bd9f3
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 1 deletion.
diff --git a/basicts/metrics/__init__.py b/basicts/metrics/__init__.py
@@ -4,8 +4,8 @@
 from .rmse import masked_rmse
 from .wape import masked_wape
 from .smape import masked_smape
-from .r_square import masked_r2
 from .corr import masked_corr
+from r_square import masked_r2
 
 ALL_METRICS = {
             'MAE': masked_mae,

diff --git a/basicts/metrics/r_square.py b/basicts/metrics/r_square.py
@@ -39,4 +39,5 @@ def masked_r2(prediction: torch.Tensor, target: torch.Tensor, null_val: float =
 
     # 计算 R^2
     loss = 1 - (ss_res / (ss_tot + eps))
+    loss = torch.nan_to_num(loss)  # Replace any NaNs in the loss with zero
     return torch.mean(loss)
diff --git a/basicts/metrics/smape.py b/basicts/metrics/smape.py
@@ -1,6 +1,7 @@
 import numpy as np
 import torch
 
+
 def masked_smape(prediction: torch.Tensor, target: torch.Tensor, null_val: float = np.nan) -> torch.Tensor:
     """
     Calculate the Masked Symmetric Mean Absolute Percentage Error (SMAPE) between predicted and target values,