From da19e0f7bcd9f257ea3a6f62d6f603ed4d261df3 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 24 Jul 2019 10:24:15 -0400
Subject: [PATCH] updated test docs

---
 tests/test_models.py | 66 +++++++++++++++++++++++---------------------
 1 file changed, 34 insertions(+), 32 deletions(-)

diff --git a/tests/test_models.py b/tests/test_models.py
index 6bac63cf..580ba1af 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -138,6 +138,40 @@ def test_multi_gpu_model_dp():
     clear_tt_dir()
 
 
+def test_amp_gpu_dp():
+    """
+    Make sure DP + AMP work
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+
+    clear_tt_dir()
+    model = get_model()
+
+    trainer = Trainer(
+        progress_bar=False,
+        experiment=get_exp(),
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        gpus=[0, 1],
+        distributed_backend='dp',
+        use_amp=True
+    )
+
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + gpu model failed to complete'
+    assert_ok_acc(trainer)
+
+    clear_tt_dir()
+
+
 def test_multi_gpu_model_ddp():
     """
     Make sure DDP works
@@ -207,38 +241,6 @@ def test_amp_gpu_ddp():
     clear_tt_dir()
 
 
-def test_amp_gpu_dp():
-    """
-    Make sure DP + AMP work
-    :return:
-    """
-    if not torch.cuda.is_available():
-        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
-        return
-
-    clear_tt_dir()
-    model = get_model()
-
-    trainer = Trainer(
-        progress_bar=False,
-        experiment=get_exp(),
-        max_nb_epochs=1,
-        train_percent_check=0.4,
-        gpus=[0, 1],
-        distributed_backend='dp',
-        use_amp=True
-    )
-
-    result = trainer.fit(model)
-
-    # correct result and ok accuracy
-    assert result == 1, 'amp + gpu model failed to complete'
-    assert_ok_acc(trainer)
-
-    clear_tt_dir()
 
 
 if __name__ == '__main__':