updated docs (#1941)

This commit is contained in:
William Falcon 2020-05-25 15:59:32 -04:00 committed by GitHub
parent 65b4352930
commit f46a7bae77
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 1 deletions

View File

@ -842,7 +842,10 @@ class Trainer(
# route to appropriate start method
# when using multi-node or DDP within a node start each module in a separate process
if self.use_ddp2:
task = int(os.environ['SLURM_LOCALID'])
if self.is_slurm_managing_tasks:
task = int(os.environ['SLURM_LOCALID'])
elif 'WORLD_SIZE' in os.environ and 'GROUP_RANK' in os.environ:
task = int(os.environ['LOCAL_RANK'])
self.ddp_train(task, model)
elif self.use_ddp:
if self.is_slurm_managing_tasks: