make sure diffusion prior trainer can operate with no warmup

Author: lucidrains

dalle2_pytorch/trainer.py

@@ -181,7 +181,7 @@ def __init__(
         eps = 1e-6,
         max_grad_norm = None,
         group_wd_params = True,
-        warmup_steps = 1,
+        warmup_steps = None,
         cosine_decay_max_steps = None,
         **kwargs
     ):
@@ -357,7 +357,8 @@ def update(self):
 
         # accelerator will ocassionally skip optimizer steps in a "dynamic loss scaling strategy"
         if not self.accelerator.optimizer_step_was_skipped:
-            with self.warmup_scheduler.dampening():
+            sched_context = self.warmup_scheduler.dampening if exists(self.warmup_scheduler) else nullcontext
+            with sched_context():
                 self.scheduler.step()
 
         if self.use_ema:

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.8.1'
+__version__ = '1.8.2'