From 0972a8e3d07a7a18e27260340676a19694c39d6a Mon Sep 17 00:00:00 2001 From: Stefaanhess Date: Tue, 14 Jul 2026 23:13:51 +0200 Subject: [PATCH] Fix EMA callback: resume training from the raw weights, not the EMA average MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Checkpoints are written during validation, when on_validation_epoch_start has copy_to()'d the EMA average into the model; restore() only runs at the next on_train_epoch_start. A run resumed from such a checkpoint therefore loads the EMA average as its model weights, and on_fit_start immediately called store(), overwriting the stored raw weights — so the subsequent restore() was a no-op and training silently continued from the EMA average instead of the raw weights a continuous run would have used. Two changes: - on_fit_start: after loading the EMA state on resume, restore the raw weights (the state's collected_params, captured by the store() at on_validation_epoch_start before the save) into the model. - on_validation_epoch_start: skip store()/copy_to() during the sanity check. on_fit_start has already stored the raw weights and copied the EMA average in, so the sanity check still validates the EMA weights; storing again would overwrite the raw weights with the average and reintroduce the bug whenever num_sanity_val_steps > 0 (the default). Fresh (non-resume) fits are unaffected. Co-Authored-By: Claude Fable 5 --- src/schnetpack/train/callbacks.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/schnetpack/train/callbacks.py b/src/schnetpack/train/callbacks.py index 10f7e13eb..edaf17996 100644 --- a/src/schnetpack/train/callbacks.py +++ b/src/schnetpack/train/callbacks.py @@ -125,6 +125,10 @@ def on_fit_start(self, trainer, pl_module: AtomisticTask): if self._to_load is not None: self.ema.load_state_dict(self._to_load) self._to_load = None + if self.ema.collected_params is not None: + # checkpoints hold the EMA average (saved during validation); + # restore the raw training weights before store()/copy_to() + self.ema.restore() # load average parameters, to have same starting point as after validation self.ema.store() @@ -141,6 +145,10 @@ def on_train_batch_end(self, trainer, pl_module: AtomisticTask, *args, **kwargs) def on_validation_epoch_start( self, trainer: "pl.Trainer", pl_module: AtomisticTask, *args, **kwargs ): + if trainer.sanity_checking: + # on_fit_start already stored the raw weights and copied the EMA + # average in; storing again would overwrite them with the average + return self.ema.store() self.ema.copy_to()