Relax dataset type check in deepspeed io (#1012)

Co-authored-by: Jeff Rasley <[email protected]>

deepspeed/runtime/dataloader.py

@@ -99,3 +99,6 @@ def _create_dataloader(self):
         self.data = (x for x in self.dataloader)
 
         return self.dataloader
+
+
+# DataLoader([(torch.randn(3, 3), torch.tensor(i % 2)) for i in range(10)], batch_size=2))

deepspeed/runtime/engine.py

@@ -845,6 +845,14 @@ def _configure_progressive_layer_drop(self):
 
         return pld
 
+    def is_map_style_dataset(obj):
+        return hasattr(obj, "__getitem__") and hasattr(obj, "__len__")
+
+    def is_iterable_style_dataset(obj):
+        return isinstance(obj,
+                          torch.utils.data.IterableDataset
+                          )  # hasattr(obj, "__iter__") should work as well
+
     def deepspeed_io(self,
                      dataset,
                      batch_size=None,
@@ -853,7 +861,7 @@ def deepspeed_io(self,
                      data_sampler=None,
                      collate_fn=None,
                      num_local_io_workers=None):
-        if not isinstance(dataset, torch.utils.data.Dataset):
+        if not (is_map_style_dataset(dataset) or is_iterable_style_dataset(dataset)):
             raise ValueError("Training data must be a torch Dataset")
 
         if data_sampler is None and (route == ROUTE_PREDICT or route == ROUTE_EVAL):