Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fixes slurm weights saving (#2339)

This commit is contained in:
William Falcon 2020-06-23 20:16:34 -04:00 committed by GitHub
parent 9446390779
commit a915280427
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
pytorch_lightning/trainer/trainer.py
View File

@ -327,8 +327,8 @@ class Trainer(
# this way we only show it on rank 0
if 'LOCAL_RANK' in os.environ:
rank_zero_only.rank = os.environ['LOCAL_RANK']
if 'SLURM_JOB_ID' in os.environ:
rank_zero_only.rank = os.environ['SLURM_JOB_ID']
if 'SLURM_LOCALID' in os.environ:
rank_zero_only.rank = os.environ['SLURM_LOCALID']
# Init callbacks
self.prepare_data_per_node = prepare_data_per_node