From 9b749fde405865f234d2f182fadf5ce5cf74b1b7 Mon Sep 17 00:00:00 2001 From: Simon Adamov Date: Sat, 25 May 2024 17:38:37 +0200 Subject: [PATCH] remove LR and bugfixing --- neural_lam/models/ar_model.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py index e496533a..29b169d4 100644 --- a/neural_lam/models/ar_model.py +++ b/neural_lam/models/ar_model.py @@ -97,9 +97,6 @@ def configure_optimizers(self): opt = torch.optim.AdamW( self.parameters(), lr=self.args.lr, betas=(0.9, 0.95) ) - if self.opt_state: - opt.load_state_dict(self.opt_state) - return opt @property @@ -598,10 +595,5 @@ def on_load_checkpoint(self, checkpoint): loaded_state_dict[new_key] = loaded_state_dict[old_key] del loaded_state_dict[old_key] if not self.restore_opt: - optimizers, lr_schedulers = self.configure_optimizers() - checkpoint["optimizer_states"] = [ - opt.state_dict() for opt in optimizers - ] - checkpoint["lr_schedulers"] = [ - sched.state_dict() for sched in lr_schedulers - ] + opt = self.configure_optimizers() + checkpoint["optimizer_states"] = [opt.state_dict()]