solved bug for energy_loss_ratio not functioning in extensive model

CederGroupHub · Feb 21, 2024 · 3f1c722 · 3f1c722
1 parent 0620c55
commit 3f1c722
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/chgnet/trainer/trainer.py b/chgnet/trainer/trainer.py
@@ -654,7 +654,9 @@ def forward(
             else:
                 e_per_atom_target = targets["e"] / prediction["atoms_per_graph"]
                 e_per_atom_pred = prediction["e"] / prediction["atoms_per_graph"]
-                out["loss"] += self.criterion(e_per_atom_target, e_per_atom_pred)
+                out["loss"] += self.energy_loss_ratio * self.criterion(
+                    e_per_atom_target, e_per_atom_pred
+                )
                 out["e_MAE"] = mae(e_per_atom_target, e_per_atom_pred)
                 out["e_MAE_size"] = prediction["e"].shape[0]