(rank: int, world_size: int, dataset_len: int)
| 46 | |
| 47 | |
| 48 | def get_local_indices(rank: int, world_size: int, dataset_len: int) -> List[int]: |
| 49 | indices = list(range(dataset_len)) |
| 50 | while len(indices) % world_size != 0: |
| 51 | indices.extend(indices[: world_size - len(indices) % world_size]) |
| 52 | indices = indices[rank::world_size] |
| 53 | return indices |
| 54 | |
| 55 | |
| 56 | def main() -> None: |