Again, why should data_mean and data_std defaults be None?

materialsvirtuallab · Aug 14, 2023 · 28cf9a3 · 28cf9a3
1 parent 3a00612
commit 28cf9a3
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/matgl/utils/training.py b/matgl/utils/training.py
@@ -144,8 +144,8 @@ class ModelLightningModule(MatglLightningModuleMixin, pl.LightningModule):
     def __init__(
         self,
         model,
-        data_mean=None,
-        data_std=None,
+        data_mean: float = 0.0,
+        data_std: float = 1.0,
         loss: str = "mse_loss",
         optimizer: Optimizer | None = None,
         scheduler=None,
@@ -177,10 +177,6 @@ def __init__(
 
         self.mae = torchmetrics.MeanAbsoluteError()
         self.rmse = torchmetrics.MeanSquaredError(squared=False)
-        if data_mean is None:
-            data_mean = torch.zeros(1)
-        if data_std is None:
-            data_std = torch.ones(1)
         self.data_mean = data_mean
         self.data_std = data_std
         self.lr = lr