Try with flipping order

src/accelerate/data_loader.py

@@ -548,8 +548,8 @@ def __iter__(self):
             batch = slice_tensors(batch, data_slice)
 
             if stop_iteration:
-                self.gradient_state._set_remainder(observed_batch_size)
                 self.gradient_state._remove_dataloader(self)
+                self.gradient_state._set_remainder(observed_batch_size)
             if batch_index >= self.skip_batches:
                 yield batch
             batch_index += 1

src/accelerate/state.py

@@ -713,6 +713,7 @@ def __init__(self):
             self.remainder = -1
             self.active_dataloader = None
             self.dataloader_references = [None]
+            self.remainder_references = [-1]
 
     @property
     def initialized(self) -> bool:
@@ -733,7 +734,6 @@ def _set_sync_gradients(self, sync_gradients):
     def _set_end_of_dataloader(self, end_of_dataloader):
         "Private function that sets whether the end of the current dataloader has been reached. Users should not have to call this."
         self.end_of_dataloader = end_of_dataloader
-        self.remainder = -1
 
     def _set_remainder(self, remainder):
         "Private function that sets the number of remaining samples at the end of the dataloader. Users should not have to call this."

src/accelerate/test_utils/scripts/test_sync.py

@@ -260,7 +260,6 @@ def main():
     if state.local_process_index == 0:
         print("**Test `accumulate` gradient accumulation with dataloader break**")
     test_dataloader_break()
-    accelerator.gradient_state._shared_state = {}
     if state.distributed_type == DistributedType.NO:
         if state.local_process_index == 0:
             print("**Test NOOP `no_sync` context manager**")