bugfix PR#343 (#351)

Co-authored-by: Abhishek Das <das.abhshk@gmail.com>
FAIR-Chem · Jun 1, 2022 · d376bcb · d376bcb
1 parent 5b6d0ac
commit d376bcb
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/ocpmodels/trainers/forces_trainer.py b/ocpmodels/trainers/forces_trainer.py
@@ -505,8 +505,10 @@ def _compute_loss(self, out, batch_list):
                         [batch.fixed.to(self.device) for batch in batch_list]
                     )
                     mask = fixed == 0
-                    if self.config["optim"]["loss_force"].startswith(
-                        "atomwise"
+                    if (
+                        self.config["optim"]
+                        .get("loss_force", "mae")
+                        .startswith("atomwise")
                     ):
                         force_mult = self.config["optim"].get(
                             "force_coefficient", 1