diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index b2f9b66e..4fe8493a 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -19,7 +19,7 @@ import tqdm from pytorch_lightning.root_module.memory import get_gpu_memory_map from pytorch_lightning.root_module.model_saving import TrainerIO from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel -from pytorch_lightning.utils.debugging import ForkedPdb +from pytorch_lightning.utils.debugging import IncompatibleArgumentsException try: from apex import amp @@ -466,7 +466,7 @@ class Trainer(TrainerIO): 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 Exception(m) + raise IncompatibleArgumentsException(m) # run through amp wrapper if self.use_amp: diff --git a/pytorch_lightning/utils/debugging.py b/pytorch_lightning/utils/debugging.py index 7a4d1445..629dda9d 100644 --- a/pytorch_lightning/utils/debugging.py +++ b/pytorch_lightning/utils/debugging.py @@ -12,4 +12,8 @@ class ForkedPdb(pdb.Pdb): sys.stdin = open('/dev/stdin') pdb.Pdb.interaction(self, *args, **kwargs) finally: - sys.stdin = _stdin \ No newline at end of file + sys.stdin = _stdin + + +class IncompatibleArgumentsException(Exception): + pass \ No newline at end of file diff --git a/tests/test_models.py b/tests/test_models.py index 96e99156..6716fef5 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -4,6 +4,7 @@ from pytorch_lightning.examples.new_project_templates.lightning_module_template from argparse import Namespace from test_tube import Experiment from pytorch_lightning.callbacks import ModelCheckpoint +from pytorch_lightning.utils.debugging import IncompatibleArgumentsException import numpy as np import warnings import torch @@ -90,17 +91,14 @@ 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 - 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 str(e) - - run_gpu_model_test(trainer_options) + trainer_options = dict( + max_nb_epochs=1, + gpus=[0, 1], + distributed_backend='dp', + use_amp=True + ) + with pytest.raises(IncompatibleArgumentsException): + run_gpu_model_test(trainer_options) def test_multi_gpu_model_ddp():