Fix batch_sampler maybe None error

src/accelerate/data_loader.py

@@ -744,7 +744,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

@@ -15,6 +15,7 @@
 import random
 import unittest
 
+import torch
 from torch.utils.data import BatchSampler, DataLoader, IterableDataset
 
 from accelerate import Accelerator
@@ -44,6 +45,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 = [
@@ -364,6 +380,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)