From a4bb80b936208b8123f46d3c4ed99dbe95393f4c Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 24 Jul 2019 19:43:38 -0400
Subject: [PATCH] dp doesnt support amp with any setting

---
 pytorch_lightning/models/trainer.py | 11 +----------
 tests/test_models.py                | 22 ----------------------
 2 files changed, 1 insertion(+), 32 deletions(-)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index e58735dc..a575a316 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -462,21 +462,12 @@ class Trainer(TrainerIO):
 
         # check for this bug (amp + dp + !01 doesn't work)
         # https://github.com/NVIDIA/apex/issues/227
-        if self.use_dp and self.use_amp and self.amp_level != 'O1':
+        if self.use_dp and self.use_amp:
             m = f'amp level {self.amp_level} with DataParallel is not supported. ' \
                 f'See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227. ' \
                 f'We recommend you switch to ddp if you want to use amp'
             raise MisconfigurationException(m)
 
-        # run through amp wrapper
-        if self.use_amp:
-
-            # An example
-            model, optimizers = amp.initialize(
-                model, self.optimizers, opt_level=self.amp_level,
-            )
-            self.optimizers = optimizers
-
         model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
 
         self.__run_pretrain_routine(model)
diff --git a/tests/test_models.py b/tests/test_models.py
index 68c97d4e..fd5fcf51 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -21,28 +21,6 @@ np.random.seed(SEED)
 # ------------------------------------------------------------------------
 # TESTS
 # ------------------------------------------------------------------------
-def test_amp_gpu_dp_ok_1():
-    """
-    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
-    model, hparams = get_model()
-    trainer_options = dict(
-        max_nb_epochs=1,
-        gpus='0, 1',  # test init with gpu string
-        distributed_backend='dp',
-        amp_level='O1',
-        use_amp=True
-    )
-    run_gpu_model_test(trainer_options, model, hparams)
-
-
 def test_early_stopping_cpu_model():
     """
     Test each of the trainer options