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

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

* empty commit

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

examples/text_to_image/train_text_to_image_sdxl.py

@@ -1252,6 +1252,10 @@ def main(args):
                 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)