Quickfix WDS2 (lucidrains#358)

Forgot to remove length as an argument to wds.WebDataset(), sorry!
ofhwei · Sep 13, 2021 · 499b4c9 · 499b4c9
1 parent 55f21ad
commit 499b4c9
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/train_dalle.py b/train_dalle.py
@@ -361,11 +361,9 @@ def tokenize(s):
     image_mapping = {
         myimg: imagepreproc
     }
-
-    num_batches = DATASET_SIZE // BATCH_SIZE
-
+
     ds = (
-        wds.WebDataset(DATASET, length=num_batches)
+        wds.WebDataset(DATASET)
         # .shuffle(is_shuffle) # Commented out for WebDataset as the behaviour cannot be predicted yet
         .map_dict(**image_text_mapping)     
         .map_dict(**image_mapping)