lower default ema beta value

dalle2_pytorch/optimizer.py

@@ -1,8 +1,10 @@
 from torch.optim import AdamW, Adam
 
 def separate_weight_decayable_params(params):
-    no_wd_params = set([param for param in params if param.ndim < 2])
+    wd_params, no_wd_params = [], []
-    wd_params = set(params) - no_wd_params
+    for param in params:
+        param_list = no_wd_params if param.ndim < 2 else wd_params
+        param_list.append(param)
     return wd_params, no_wd_params
 
 def get_optimizer(
@@ -25,8 +27,8 @@ def get_optimizer(
         wd_params, no_wd_params = separate_weight_decayable_params(params)
 
         params = [
-            {'params': list(wd_params)},
+            {'params': wd_params},
-            {'params': list(no_wd_params), 'weight_decay': 0},
+            {'params': no_wd_params, 'weight_decay': 0},
         ]
 
     return AdamW(params, lr = lr, weight_decay = wd, betas = betas, eps = eps)

dalle2_pytorch/trainer.py

@@ -178,7 +178,7 @@ class EMA(nn.Module):
     def __init__(
         self,
         model,
-        beta = 0.9999,
+        beta = 0.99,
         update_after_step = 1000,
         update_every = 10,
     ):

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.6.2'
+__version__ = '0.6.5'

setup.py

@@ -32,6 +32,7 @@ setup(
     'embedding-reader',
     'kornia>=0.5.4',
     'numpy',
+    'packaging',
     'pillow',
     'pydantic',
     'resize-right>=0.0.2',