From a8d32d73d234038dddc8ebf3779c7723421757f8 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Tue, 23 Jun 2020 20:15:49 -0400 Subject: [PATCH] fixes slurm weights saving --- 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 c207d8255d0c2..2c4a47b47ad5e 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -327,8 +327,8 @@ def __init__( # 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