From f743ccf7ef36632a12e0bf2d658b7bd9ea325fb6 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Wed, 24 Jun 2026 09:13:01 -0600
Subject: [PATCH] Save before last sample

---
 jobs/process/BaseSDTrainProcess.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py
index 3096c752..e3817a38 100644
--- a/jobs/process/BaseSDTrainProcess.py
+++ b/jobs/process/BaseSDTrainProcess.py
@@ -2607,12 +2607,13 @@ class BaseSDTrainProcess(BaseTrainProcess):
             self.progress_bar.close()
         if self.train_config.free_u:
             self.sd.pipeline.disable_freeu()
+        if self.accelerator.is_main_process:
+            self.save()
         if not self.train_config.disable_sampling:
             self.sample(self.step_num)
             self.logger.commit(step=self.step_num)
         print_acc("")
         if self.accelerator.is_main_process:
-            self.save()
             self.logger.finish()
         self.accelerator.end_training()