Fix batch_sampler maybe None error (#3025)

* Fix batch_sampler maybe None

For more details, see: #3011

* Update test_data_loader.py

Add unit test for dataloader with batch_size=None when using Iterabledataset

* Update tests/test_data_loader.py

Co-authored-by: Zach Mueller <[email protected]>

* Fix inconsistent indentation

Fix inconsistent indentation

---------

Co-authored-by: Zach Mueller <[email protected]>

src/accelerate/data_loader.py

@@ -820,7 +820,7 @@ def set_epoch(self, epoch: int):
         # In case it is manually passed in, the user can set it to what they like
         if self.iteration != epoch:
             self.iteration = epoch
-        if hasattr(self.batch_sampler.sampler, "set_epoch"):
+        if hasattr(self.batch_sampler, "sampler") and hasattr(self.batch_sampler.sampler, "set_epoch"):
             self.batch_sampler.sampler.set_epoch(epoch)
         elif hasattr(self.dataset, "set_epoch"):
             self.dataset.set_epoch(epoch)

tests/test_data_loader.py

@@ -64,6 +64,21 @@ def __iter__(self):
             stop = random.random() < self.p_stop
 
 
+class SimpleIterableDataset(IterableDataset):
+    def __init__(self, num_samples=1000):
+        self.num_samples = num_samples
+
+    def __iter__(self):
+        for _ in range(self.num_samples):
+            yield torch.rand(1)
+
+    def __len__(self):
+        return self.num_samples
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+
+
 class DataLoaderTester(unittest.TestCase):
     def check_batch_sampler_shards(self, batch_sampler, expected, split_batches=False, even_batches=True):
         batch_sampler_shards = [
@@ -384,6 +399,14 @@ def test_iterable_dataset_shard(self):
         self.check_iterable_dataset_shards(dataset, seed, batch_size=4, drop_last=False, split_batches=True)
         self.check_iterable_dataset_shards(dataset, seed, batch_size=4, drop_last=True, split_batches=True)
 
+    def test_iterable_dataset_using_none_batch_size(self):
+        dataset = SimpleIterableDataset(100)
+        accelerator = Accelerator()
+        dataloader = DataLoader(dataset, batch_size=None)
+        dataloader = accelerator.prepare(dataloader)
+        for d in dataloader:
+            assert isinstance(d, torch.Tensor)
+
     def test_skip_batch_sampler(self):
         batch_sampler = BatchSampler(range(16), batch_size=4, drop_last=False)
         new_batch_sampler = SkipBatchSampler(batch_sampler, 2)