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 = device))
+        self.register_buffer('step', torch.tensor([0], device = self.device))
 
     # accelerator wrappers
 

dalle2_pytorch/version.py

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