From d3ca8360244e7757b631c166b1ed2a437dd6d0d5 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Sat, 5 Oct 2019 16:05:42 -0400 Subject: [PATCH] cleaning up demos --- pytorch_lightning/trainer/trainer.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index 9969975b..ee6c24af 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -704,7 +704,7 @@ class Trainer(TrainerIO): warnings.warn(msg) break - if self.use_ddp: + if self.use_ddp or self.use_ddp2: # wait for all processes to catch up dist.barrier() @@ -718,13 +718,12 @@ class Trainer(TrainerIO): # ----------------------------- def fit(self, model): # when using multi-node or DDP within a node start each module in a separate process - if self.use_ddp: + if self.use_ddp2: + task = int(os.environ['SLURM_LOCALID']) + self.ddp_train(task, model) - if self.use_ddp2: - task = int(os.environ['SLURM_LOCALID']) - self.ddp_train(task, model) - - elif self.is_slurm_managing_tasks: + elif self.use_ddp: + if self.is_slurm_managing_tasks: task = int(os.environ['SLURM_LOCALID']) self.ddp_train(task, model) else: