switched cpu amp order

2019-07-25 11:10:21 -04:00 · 2019-07-25 11:10:21 -04:00 · aadf8e16aa
parent 4b04dc06d4
commit aadf8e16aa
1 changed files with 2 additions and 0 deletions
--- 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)