Merge pull request #230 from mehdidc/data_sampler_set_epoch

use data_sampler.set_epoch to avoid having the same order in each epoch

Author: lucidrains

train_dalle.py

@@ -312,6 +312,8 @@ def save_model(path):
 # training
 
 for epoch in range(EPOCHS):
+    if data_sampler:
+        data_sampler.set_epoch(epoch)
     for i, (text, images) in enumerate(distr_dl):
         if args.fp16:
             images = images.half()