bugfixed: stuck when training with dist_train.sh, support tcp_port (o…

…pen-mmlab#784)
serkankartal · Feb 2, 2022 · 7ce6a2b · 7ce6a2b
1 parent 274c90c
commit 7ce6a2b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/pcdet/utils/common_utils.py b/pcdet/utils/common_utils.py
@@ -161,9 +161,11 @@ def init_dist_slurm(tcp_port, local_rank, backend='nccl'):
 def init_dist_pytorch(tcp_port, local_rank, backend='nccl'):
     if mp.get_start_method(allow_none=True) is None:
         mp.set_start_method('spawn')
-
+    os.environ['MASTER_PORT'] = str(tcp_port)
+    os.environ['MASTER_ADDR'] = 'localhost'
     num_gpus = torch.cuda.device_count()
     torch.cuda.set_device(local_rank % num_gpus)
+
     dist.init_process_group(
         backend=backend,
         # init_method='tcp://127.0.0.1:%d' % tcp_port,