set ability to do warmup steps for each unet during training

dalle2_pytorch/trainer.py

@@ -550,7 +550,7 @@ class DecoderTrainer(nn.Module):
         if only_model:
             return loaded_obj
 
-        for ind, last_step in zip(range(0, self.num_unets), self.steps.cpu().unbind()):
+        for ind, last_step in zip(range(0, self.num_unets), self.steps.tolist()):
 
             optimizer_key = f'optim{ind}'
             optimizer = getattr(self, optimizer_key)

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.16.5'
+__version__ = '0.16.6'