From a915280427ed02b5cae386d0e32792088fa7b389 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Tue, 23 Jun 2020 20:16:34 -0400 Subject: [PATCH] fixes slurm weights saving (#2339) --- pytorch_lightning/trainer/trainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index c207d8255d..2c4a47b47a 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -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