Add missing restore() EMA call in train SDXL script (huggingface#7599)

* Restore unet params back to normal from EMA when validation call is finished

* empty commit

---------

Co-authored-by: Sayak Paul <[email protected]>

Author: christopher-beckham

examples/text_to_image/train_text_to_image_sdxl.py

@@ -1252,6 +1252,10 @@ def compute_time_ids(original_size, crops_coords_top_left):
                 del pipeline
                 torch.cuda.empty_cache()
 
+                if args.use_ema:
+                    # Switch back to the original UNet parameters.
+                    ema_unet.restore(unet.parameters())
+
     accelerator.wait_for_everyone()
     if accelerator.is_main_process:
         unet = unwrap_model(unet)