Fix final checkpoint bug (#433)

src/metatrain/experimental/nanopet/trainer.py

@@ -259,6 +259,7 @@ def systems_and_targets_to_dtype(
         if self.best_loss is None:
             self.best_loss = float("inf")
         logger.info("Starting training")
+        epoch = start_epoch
         for epoch in range(start_epoch, start_epoch + self.hypers["num_epochs"]):
             if is_distributed:
                 sampler.set_epoch(epoch)
@@ -314,8 +315,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 +460,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()
 

src/metatrain/experimental/soap_bpnn/trainer.py

@@ -255,6 +255,7 @@ def train(
         if self.best_loss is None:
             self.best_loss = float("inf")
         logger.info("Starting training")
+        epoch = start_epoch
         for epoch in range(start_epoch, start_epoch + self.hypers["num_epochs"]):
             if is_distributed:
                 sampler.set_epoch(epoch)
@@ -449,6 +450,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()