diff --git a/train.py b/train.py index c3dae69ac025e44b1c95b5b281fcd38ae3b2dfbb..d95e5e6c873b606aec720423051e10e6555aa6e6 100644 --- a/train.py +++ b/train.py @@ -563,8 +563,7 @@ if __name__ == '__main__': world_rank = 0 local_rank = 0 if params['world_size'] > 1: - dist.init_process_group(backend='nccl', - init_method='env://', timeout=timedelta(seconds=120)) + dist.init_process_group(init_method='env://', timeout=timedelta(seconds=120)) local_rank = int(os.environ["LOCAL_RANK"]) args.gpu = local_rank world_rank = dist.get_rank()