Why should weights be None. A 0.0 weight is basically the same.

matgl/utils/training.py

@@ -257,8 +257,8 @@ def __init__(
         element_refs: np.ndarray | None = None,
         energy_weight: float = 1.0,
         force_weight: float = 1.0,
-        stress_weight: float | None = None,
-        site_wise_weight: float | None = None,
+        stress_weight: float = 0.0,
+        site_wise_weight: float = 0.0,
         data_mean: float = 0.0,
         data_std: float = 1.0,
         calc_stress: bool = False,
@@ -308,7 +308,7 @@ def __init__(
         self.decay_steps = decay_steps
         self.decay_alpha = decay_alpha
 
-        calc_site_wise = site_wise_weight is not None
+        calc_site_wise = site_wise_weight != 0
         self.model = Potential(
             model=model,
             element_refs=element_refs,
@@ -369,10 +369,6 @@ def step(self, batch: tuple):
             loss=self.loss,  # type: ignore
             preds=preds,
             labels=labels,
-            energy_weight=self.energy_weight,
-            force_weight=self.force_weight,
-            stress_weight=self.stress_weight,
-            site_wise_weight=self.site_wise_weight,
             num_atoms=num_atoms,
         )
         batch_size = preds[0].numel()
@@ -384,10 +380,6 @@ def loss_fn(
         loss: nn.Module,
         labels: tuple,
         preds: tuple,
-        energy_weight: float | None = None,
-        force_weight: float | None = None,
-        stress_weight: float | None = None,
-        site_wise_weight: float | None = None,
         num_atoms: int | None = None,
     ):
         """Compute losses for EFS.
@@ -396,10 +388,6 @@ def loss_fn(
             loss: Loss function.
             labels: Labels.
             preds: Predictions
-            energy_weight: Weight for energy loss.
-            force_weight: Weight for force loss.
-            stress_weight: Weight for stress loss.
-            site_wise_weight: Weight for site-wise loss.
             num_atoms: Number of atoms.
 
         Returns::
@@ -431,19 +419,19 @@ def loss_fn(
         m_mae = torch.zeros(1)
         m_rmse = torch.zeros(1)
 
-        total_loss = energy_weight * e_loss + force_weight * f_loss
+        total_loss = self.energy_weight * e_loss + self.force_weight * f_loss
 
-        if stress_weight is not None:
+        if self.stress_weight:
             s_loss = loss(labels[2], preds[2])
             s_mae = self.mae(labels[2], preds[2])
             s_rmse = self.rmse(labels[2], preds[2])
-            total_loss = total_loss + stress_weight * s_loss
+            total_loss = total_loss + self.stress_weight * s_loss
 
-        if site_wise_weight is not None:
+        if self.site_wise_weight:
             m_loss = loss(labels[3], preds[3])
             m_mae = self.mae(labels[3], preds[3])
             m_rmse = self.rmse(labels[3], preds[3])
-            total_loss = total_loss + site_wise_weight * m_loss
+            total_loss = total_loss + self.site_wise_weight * m_loss
 
         return {
             "Total_Loss": total_loss,

tests/utils/test_training.py

@@ -119,7 +119,6 @@ def test_m3gnet_training(self, LiFePO4, BaNiO3):
 
         trainer.fit(model=lit_model, train_dataloaders=train_loader, val_dataloaders=val_loader)
 
-        model = model.to(torch.device(device))
         pred_LFP_energy = model.predict_structure(LiFePO4)
         pred_BNO_energy = model.predict_structure(BaNiO3)
 
@@ -167,7 +166,6 @@ def test_m3gnet_property_training(self, LiFePO4, BaNiO3):
 
         trainer.fit(model=lit_model, train_dataloaders=train_loader, val_dataloaders=val_loader)
 
-        model = model.to(torch.device(device))
         pred_LFP_energy = model.predict_structure(LiFePO4)
         pred_BNO_energy = model.predict_structure(BaNiO3)