Fix final checkpoint bug

metatensor · Dec 17, 2024 · a6f0dad · a6f0dad
1 parent 8ae4088
commit a6f0dad
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/metatrain/experimental/nanopet/trainer.py b/src/metatrain/experimental/nanopet/trainer.py
@@ -314,8 +314,6 @@ def systems_and_targets_to_dtype(
                 if self.hypers["log_mae"]:
                     train_mae_calculator.update(predictions, targets)
 
-                # count += 1
-
             finalized_train_info = train_rmse_calculator.finalize(
                 not_per_atom=["positions_gradients"] + per_structure_targets,
                 is_distributed=is_distributed,
@@ -461,6 +459,7 @@ def systems_and_targets_to_dtype(
                     )
 
         # prepare for the checkpoint that will be saved outside the function
+        self.epoch = epoch
         self.optimizer_state_dict = optimizer.state_dict()
         self.scheduler_state_dict = lr_scheduler.state_dict()
 

diff --git a/src/metatrain/experimental/soap_bpnn/trainer.py b/src/metatrain/experimental/soap_bpnn/trainer.py
@@ -449,6 +449,7 @@ def train(
                     )
 
         # prepare for the checkpoint that will be saved outside the function
+        self.epoch = epoch
         self.optimizer_state_dict = optimizer.state_dict()
         self.scheduler_state_dict = lr_scheduler.state_dict()