From aadf8e16aa56c81bf3f840b7d200a3c7f1a904f2 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Thu, 25 Jul 2019 11:10:21 -0400
Subject: [PATCH] switched cpu amp order

---
 tests/test_models.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/test_models.py b/tests/test_models.py
index 482fcae750..ea0741da41 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -55,7 +55,9 @@ def test_amp_gpu_ddp_slurm_managed():
         warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
         return
 
+    # simulate setting slurm flags
     os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+    os.environ['SLURM_LOCALID'] = 0
 
     hparams = get_hparams()
     model = LightningTestModel(hparams)