remove deprecated averaged_model

teticio · Feb 3, 2023 · 7bd9ee1 · 7bd9ee1
1 parent d8c6a4b
commit 7bd9ee1
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/audiodiffusion/__init__.py b/audiodiffusion/__init__.py
@@ -9,7 +9,7 @@
 # from diffusers import AudioDiffusionPipeline
 from .pipeline_audio_diffusion import AudioDiffusionPipeline
 
-VERSION = "1.5.0"
+VERSION = "1.5.1"
 
 
 class AudioDiffusion:

diff --git a/scripts/train_unet.py b/scripts/train_unet.py
@@ -304,10 +304,12 @@ def transforms(examples):
             if ((epoch + 1) % args.save_model_epochs == 0
                     or (epoch + 1) % args.save_images_epochs == 0
                     or epoch == args.num_epochs - 1):
+                unet = accelerator.unwrap_model(model)
+                if args.use_ema:
+                    ema_model.copy_to(unet.parameters())
                 pipeline = AudioDiffusionPipeline(
                     vqvae=vqvae,
-                    unet=accelerator.unwrap_model(
-                        ema_model.averaged_model if args.use_ema else model),
+                    unet=unet,
                     mel=mel,
                     scheduler=noise_scheduler,
                 )