pass parameter for Trainer to EMA properly

andrewmoise · Jun 21, 2022 · 75ea49a · 75ea49a
1 parent 8c3609a
commit 75ea49a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
@@ -590,7 +590,7 @@ def __init__(
         gradient_accumulate_every = 2,
         amp = False,
         step_start_ema = 2000,
-        update_ema_every = 10,
+        ema_update_every = 10,
         save_and_sample_every = 1000,
         results_folder = './results',
         augment_horizontal_flip = True
@@ -599,8 +599,7 @@ def __init__(
         self.image_size = diffusion_model.image_size
 
         self.model = diffusion_model
-        self.ema = EMA(diffusion_model, beta = ema_decay)
-        self.update_ema_every = update_ema_every
+        self.ema = EMA(diffusion_model, beta = ema_decay, update_every = ema_update_every)
 
         self.step_start_ema = step_start_ema
         self.save_and_sample_every = save_and_sample_every

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.21.0',
+  version = '0.21.1',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',