From 89ff04cfe2735a28d74141c2e65448bd84619cd9 Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Mon, 16 May 2022 11:54:34 -0700 Subject: [PATCH] final tweak to EMA class --- dalle2_pytorch/trainer.py | 6 +++++- setup.py | 2 +- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/dalle2_pytorch/trainer.py b/dalle2_pytorch/trainer.py index 594e389..fdc9b0f 100644 --- a/dalle2_pytorch/trainer.py +++ b/dalle2_pytorch/trainer.py @@ -195,7 +195,11 @@ class EMA(nn.Module): def update(self): self.step += 1 - if self.step <= self.update_after_step or (self.step % self.update_every) != 0: + if (self.step % self.update_every) != 0: + return + + if self.step <= self.update_after_step: + self.copy_params_from_model_to_ema() return if not self.initted: diff --git a/setup.py b/setup.py index 930e90e..3943dd4 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,7 @@ setup( 'dream = dalle2_pytorch.cli:dream' ], }, - version = '0.2.43', + version = '0.2.44', license='MIT', description = 'DALL-E 2', author = 'Phil Wang',