Merge pull request #449 from ACEsuit/multihead_validation_print

validation loss printed during training is not actually for the specific head
ACEsuit · Jun 6, 2024 · 14ec4ec · 14ec4ec
2 parents a92dba2 + 9616a41
commit 14ec4ec
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,