[fix]fix dataloader: honor min_batch_size when device_mesh is None

src/twinkle/dataloader/dataloader.py

@@ -134,6 +134,7 @@ def __iter__(self):
                 _iter._dataset_fetcher.drop_last,
                 self.batch_size,
                 self.device_mesh,
+                min_batch_size=self.min_batch_size,
                 max_retries=self.max_retries)
         return self._tracking_iter(_iter)
 

src/twinkle/dataloader/device_mesh_fetcher.py

@@ -77,9 +77,8 @@ def fetch(self, _):
         else:
             data = next(self.dataset_iter)
 
+        if self.min_batch_size is not None and len(data) < self.min_batch_size:
+            raise StopIteration
         if self.device_mesh:
-            if len(data) < self.min_batch_size:
-                raise StopIteration
-            else:
-                data = data[self.device_mesh.get_slice(len(data))]
+            data = data[self.device_mesh.get_slice(len(data))]
         return self.collate_fn(data)

src/twinkle/dataloader/device_mesh_sampler.py

@@ -34,13 +34,12 @@ def __iter__(self):
                 batch = batch[self.skip_samples - skipped:]
                 skipped = self.skip_samples
 
+            if self.min_batch_size is not None and len(batch) < self.min_batch_size:
+                return
             if not self.device_mesh:
                 yield batch
             else:
-                if len(batch) < self.min_batch_size:
-                    return
-                else:
-                    yield batch[self.device_mesh.get_slice(len(batch))]
+                yield batch[self.device_mesh.get_slice(len(batch))]
 
     def __len__(self):
         return len(self.original_sampler)