default the device to the device that the diffusion prior parameters are on, if the trainer was never given the accelerator nor device

dalle2_pytorch/trainer.py

@@ -228,7 +228,7 @@ class DiffusionPriorTrainer(nn.Module):
 
         # track steps internally
 
-        self.register_buffer('step', torch.tensor([0], device = self.device))
+        self.register_buffer('step', torch.tensor([0], device = device))
 
     # accelerator wrappers
 
@@ -473,7 +473,7 @@ class DecoderTrainer(nn.Module):
 
         lr, wd, eps, warmup_steps = map(partial(cast_tuple, length = self.num_unets), (lr, wd, eps, warmup_steps))
 
-        assert all([unet_lr <= 1e-2 for unet_lr in lr]), 'your learning rate is too high, recommend sticking with 1e-4, at most 5e-4'
+        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 = []
         schedulers = []

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.16.13'
+__version__ = '0.16.11'