Fix immediate StopIteration in DataLoaderShard

src/accelerate/data_loader.py

@@ -558,30 +558,34 @@ def __iter__(self):
 
         self.set_epoch(self.iteration)
         dataloader_iter = self.base_dataloader.__iter__()
+
         # We iterate one batch ahead to check when we are at the end
         try:
-            current_batch = next(dataloader_iter)
+            next_batch = next(dataloader_iter)
         except StopIteration:
-            yield
+            self.end_of_dataloader = True
+            # Need to call update to set "_iterator_finished"
+            self._update_state_dict()
 
         batch_index = 0
-        while True:
+        while not self.end_of_dataloader:
+            current_batch = next_batch
+            # But we still move it to the device so it is done before `StopIteration` is reached
+            if self.device is not None:
+                current_batch = send_to_device(current_batch, self.device, non_blocking=self._non_blocking)
+
+            # We need to update the state dict before iterating again
+            self._update_state_dict()
             try:
-                # But we still move it to the device so it is done before `StopIteration` is reached
-                if self.device is not None:
-                    current_batch = send_to_device(current_batch, self.device, non_blocking=self._non_blocking)
-                self._update_state_dict()
                 next_batch = next(dataloader_iter)
-                if batch_index >= self.skip_batches:
-                    yield current_batch
-                batch_index += 1
-                current_batch = next_batch
             except StopIteration:
                 self.end_of_dataloader = True
+                # Need to call update again to set "_iterator_finished"
                 self._update_state_dict()
-                if batch_index >= self.skip_batches:
-                    yield current_batch
-                break
+
+            if batch_index >= self.skip_batches:
+                yield current_batch
+            batch_index += 1
 
         self.iteration += 1
         self.end()