just make sure decoder learning rate is reasonable and help out budding researchers

dalle2_pytorch/trainer.py

@@ -451,6 +451,8 @@ class DecoderTrainer(nn.Module):
 
         lr, wd, eps = map(partial(cast_tuple, length = self.num_unets), (lr, wd, eps))
 
+        assert all([unet_lr < 1e-3 for unet_lr in lr]), 'your learning rate is too high, recommend sticking with 1e-4, at most 5e-4'
+
         optimizers = []
 
         for unet, unet_lr, unet_wd, unet_eps in zip(decoder.unets, lr, wd, eps):

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.11.4'
+__version__ = '0.11.5'

train_decoder.py

@@ -258,8 +258,8 @@ def train(
     is_master = accelerator.process_index == 0
 
     trainer = DecoderTrainer(
-        accelerator,
-        decoder,
+        decoder=decoder,
+        accelerator=accelerator,
         **kwargs
     )