Do not try to wrap WebDatasets with DeepSpeed (lucidrains#367)

Wrapping causes errors due to PyTorch's
`torch.data.utils.DistributedSampler` not being applicable to
`torch.data.utils.IterableDataset`s (which WebDatasets are
implementing).

Fix lucidrains#359.

Author: janEbert

train_dalle.py

@@ -492,7 +492,11 @@ def filter_dataset(item): # For e.g. C@H which (rarely) has no caption available
     model=dalle,
     optimizer=opt,
     model_parameters=get_trainable_params(dalle),
-    training_data=ds if using_deepspeed else dl,
+    training_data=(
+        (None if ENABLE_WEBDATASET else ds)
+        if using_deepspeed
+        else dl
+    ),
     # Do not pass the LR scheduler to DeepSpeed so we can manually
     # advance it.
     lr_scheduler=scheduler if LR_DECAY and not using_deepspeed else None,