r"""Batch method provides a batch indices generator.
(self)
| 126 | return indices |
| 127 | |
| 128 | def batch(self) -> Iterator[List[Any]]: |
| 129 | r"""Batch method provides a batch indices generator.""" |
| 130 | indices = list(self.sample()) |
| 131 | |
| 132 | # user might pass the world_size parameter without dist, |
| 133 | # so dist.is_distributed() should not be used |
| 134 | if self.world_size > 1: |
| 135 | indices = self.scatter(indices) |
| 136 | |
| 137 | batch = [] |
| 138 | for idx in indices: |
| 139 | batch.append(idx) |
| 140 | if len(batch) == self.batch_size: |
| 141 | yield batch |
| 142 | batch = [] |
| 143 | |
| 144 | if len(batch) > 0 and not self.drop_last: |
| 145 | yield batch |
| 146 | |
| 147 | |
| 148 | class StreamSampler(Sampler): |