From 7045a01375cf3f3b078f77ba04b4d1a8f967fb4e Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 21 May 2025 02:27:55 -0600 Subject: [PATCH] Fixed issue saving optimizer in some instances. --- jobs/process/BaseSDTrainProcess.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index 393ba831..b4f768d9 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -629,7 +629,10 @@ class BaseSDTrainProcess(BaseTrainProcess): try: filename = f'optimizer.pt' file_path = os.path.join(self.save_root, filename) - state_dict = unwrap_model(self.optimizer).state_dict() + try: + state_dict = unwrap_model(self.optimizer).state_dict() + except Exception as e: + state_dict = self.optimizer.state_dict() torch.save(state_dict, file_path) print_acc(f"Saved optimizer to {file_path}") except Exception as e: