ema: fix saving of inference-compatible weights to reference the corr…

…ect parameters instead of saving a 2nd copy of the main model

helpers/training/save_hooks.py

@@ -210,7 +210,7 @@ def _save_lora(self, models, weights, output_dir):
             logger.info("Saving EMA model to disk.")
             trainable_parameters = [
                 p
-                for p in self.ema_model.parameters()
+                for p in self._primary_model().parameters()
                 if p.requires_grad
             ]
             self.ema_model.store(trainable_parameters)