r"""Initialize distributed training
(local_rank, backend='nccl', **kwargs)
| 4 | |
| 5 | |
| 6 | def init_dist(local_rank, backend='nccl', **kwargs): |
| 7 | r"""Initialize distributed training""" |
| 8 | if dist.is_available(): |
| 9 | if dist.is_initialized(): |
| 10 | return torch.cuda.current_device() |
| 11 | torch.cuda.set_device(local_rank) |
| 12 | dist.init_process_group(backend=backend, init_method='env://', **kwargs) |
| 13 | |
| 14 | |
| 15 | def get_rank(): |