fixed nccl init

2019-07-12 15:59:12 -04:00 · 2019-07-12 15:59:12 -04:00 · 4f5eef2e78
parent 6c02afefca
commit 4f5eef2e78
1 changed files with 2 additions and 1 deletions
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@ -386,7 +386,8 @@ class Trainer(TrainerIO):
        root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
        os.environ['MASTER_ADDR'] = root_node
        os.environ['MASTER_PORT'] = f'{port}'
-        dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
+        # dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
+        dist.init_process_group("nccl")

    def __run_pretrain_routine(self, model):
        """