When printing validation loss during training, pass loss for head, no…

…t partial sum over heads calculated so far
ACEsuit · Jun 6, 2024 · 9616a41 · 9616a41
1 parent a92dba2
commit 9616a41
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mace/tools/train.py b/mace/tools/train.py
@@ -161,7 +161,7 @@ def train(
         )
         valid_loss += valid_loss_head
         valid_err_log(
-            valid_loss, eval_metrics, logger, log_errors, None, valid_loader_name
+            valid_loss_head, eval_metrics, logger, log_errors, None, valid_loader_name
         )
 
     while epoch < max_num_epochs:
@@ -224,7 +224,7 @@ def train(
                     )
                     valid_loss += valid_loss_head
                     valid_err_log(
-                        valid_loss,
+                        valid_loss_head,
                         eval_metrics,
                         logger,
                         log_errors,