modified: src/trainer/utils.py

Ruyi-Zha · Sep 18, 2023 · 6811fd9 · 6811fd9
1 parent 9af16c6
commit 6811fd9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/trainer/utils.py b/src/trainer/utils.py
@@ -365,7 +365,7 @@ def cal_rmse(a, b, mask):
     if len(mask.shape) == len(a.shape) - 1:
         mask = mask[..., None]
     mask_sum = np.sum(mask) + 1e-10
-    rmse = (((a - b)**2 * mask).sum() / (mask_sum * 3.0))**0.5
+    rmse = (((a - b)**2 * mask).sum() / (mask_sum))**0.5
     return rmse