From 75ea49a7eff6e5b5a922f0fa2a9a2b06eab137be Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Tue, 21 Jun 2022 07:37:43 -0700 Subject: [PATCH] pass parameter for Trainer to EMA properly --- denoising_diffusion_pytorch/denoising_diffusion_pytorch.py | 5 ++--- setup.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index 96222c5..3bb5435 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py @@ -590,7 +590,7 @@ class Trainer(object): 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 @@ class Trainer(object): 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 index e021727..dbf3bc6 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ from setuptools import setup, find_packages 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',