From 08f3605e5a70d61289f9079351fe62e0c7daacac Mon Sep 17 00:00:00 2001 From: MJordahn Date: Tue, 2 Jun 2026 22:14:03 +0100 Subject: [PATCH] Changed TrainCheckpointCallback to always save on_train_end if a TrainCheckpointCallback is passed to training config --- src/fairchem/core/components/train/train_runner.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/src/fairchem/core/components/train/train_runner.py b/src/fairchem/core/components/train/train_runner.py index 78b49e829d..be6108a476 100644 --- a/src/fairchem/core/components/train/train_runner.py +++ b/src/fairchem/core/components/train/train_runner.py @@ -111,12 +111,11 @@ def on_train_step_start(self, state: State, unit: TTrainUnit) -> None: shutil.rmtree(dir) def on_train_end(self, state: State, unit: TTrainUnit) -> None: - if self.checkpoint_every_n_steps is not None: - # also always checkpoint on train end - assert ( - self.save_callback - ), "Must initialize set_checkpoint_call_backs from Runner!" - self.save_callback(os.path.join(self.checkpoint_dir, "final")) + # always checkpoint on train end + assert ( + self.save_callback + ), "Must initialize set_checkpoint_call_backs from Runner!" + self.save_callback(os.path.join(self.checkpoint_dir, "final")) class TrainEvalRunner(Runner):