aljaz/DALLE2-pytorch
1
1
Fork 0
mirror of https://github.com/lucidrains/DALLE2-pytorch.git synced 2025-12-19 17:54:20 +01:00
Code Issues Packages Projects Releases Wiki Activity

final tweak to EMA class

Browse Source
This commit is contained in:
Phil Wang
2022-05-16 11:54:34 -07:00
parent f4016f6302
commit 89ff04cfe2
2 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
dalle2_pytorch/trainer.py
View File

@@ -195,7 +195,11 @@ class EMA(nn.Module):
def update(self): def update(self):
self.step += 1 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 return
if not self.initted: if not self.initted:

2
setup.py
View File

@@ -10,7 +10,7 @@ setup(
'dream = dalle2_pytorch.cli:dream' 'dream = dalle2_pytorch.cli:dream'
], ],
}, },
version = '0.2.43', version = '0.2.44',
license='MIT', license='MIT',
description = 'DALL-E 2', description = 'DALL-E 2',
author = 'Phil Wang', author = 'Phil Wang',