From 4d559d9e3b5badc5a84f4cf1e83026b80c981e31 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 24 Jul 2019 14:12:41 -0400
Subject: [PATCH] fixed amp bug

---
 tests/test_models.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/tests/test_models.py b/tests/test_models.py
index b9bb4c65..f9b7b91c 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -98,13 +98,15 @@ def test_amp_gpu_dp():
         warnings.warn('test_amp_gpu_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
         return
 
-    trainer_options = dict(
-        progress_bar=False,
-        max_nb_epochs=1,
-        gpus=[0, 1],
-        distributed_backend='dp',
-        use_amp=True
-    )
+    try:
+        trainer_options = dict(
+            max_nb_epochs=1,
+            gpus=[0, 1],
+            distributed_backend='dp',
+            use_amp=True
+        )
+    except Exception as e:
+        assert 'https://github.com/NVIDIA/apex/issues/227' in e
 
     run_gpu_model_test(trainer_options)