From 82190b41e6fb6e90de6823e091f3accabfea552b Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Sun, 7 Jun 2026 11:11:49 -0600
Subject: [PATCH] Fix bug where EMA was not being initialized even when setup
 in the config. EMA will now properly be setup and used.

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

diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py
index 2ddf684..6bd6591 100644
--- a/jobs/process/BaseSDTrainProcess.py
+++ b/jobs/process/BaseSDTrainProcess.py
@@ -2026,6 +2026,9 @@ class BaseSDTrainProcess(BaseTrainProcess):
             # Update the learning rates if they changed
             # optimizer.param_groups = previous_params
 
+        # set up the ema now that the optimizer (and its params) are ready
+        self.setup_ema()
+
         lr_scheduler_params = self.train_config.lr_scheduler_params
 
         # make sure it had bare minimum