Fix data loader

borchero · Dec 28, 2021 · 48288dd · 48288dd
1 parent 82895c4
commit 48288dd
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/lightkit/data/loader.py b/lightkit/data/loader.py
@@ -49,10 +49,20 @@ def __init__(self, dataset: Dataset[T_co], **kwargs: Any):
             kwargs.setdefault("collate_fn", collate_tuple)
 
         super().__init__(dataset, **kwargs)  # type: ignore
+        self.custom_batching = self.num_workers == 0 and (
+            isinstance(self.batch_sampler, RangeBatchSampler)
+            or (
+                self.batch_sampler is not None
+                and hasattr(self.batch_sampler, "sampler")
+                and isinstance(self.batch_sampler.sampler, RangeBatchSampler)
+            )
+        )
 
     def __iter__(self) -> Iterator[Any]:  # pylint: disable=inconsistent-return-statements
-        if not (isinstance(self.dataset, TensorDataset) and self.num_workers == 0):
-            return super().__iter__()
+        if not self.custom_batching:
+            for item in super().__iter__():
+                yield item
+            return
 
         for indices in self.batch_sampler:
             if isinstance(indices, range):