Prior updates (#211)

* update configs for prior

add prior warmup to config

update example prior config

* update prior trainer & script

add deepspeed amp & warmup

adopt full accelerator support

reload at sample point

finish epoch resume code

* update tracker save method for prior

* helper functions for prior_loader

dalle2_pytorch/dataloaders/prior_loader.py

@@ -67,6 +67,15 @@ class PriorEmbeddingDataset(IterableDataset):
     def __str__(self):
         return f"<PriorEmbeddingDataset: start: {self.start}, stop: {self.stop}, len: {self.__len__()}>"
 
+    def set_start(self, start):
+        """
+        Adjust the starting point within the reader, useful for resuming an epoch
+        """
+        self.start = start
+
+    def get_start(self):
+        return self.start
+
     def get_sample(self):
         """
         pre-proocess data from either reader into a common format