diff --git a/pytorch_lightning/trainer/dp_mixin.py b/pytorch_lightning/trainer/dp_mixin.py index f46d550d..6bc08704 100644 --- a/pytorch_lightning/trainer/dp_mixin.py +++ b/pytorch_lightning/trainer/dp_mixin.py @@ -87,12 +87,15 @@ class TrainerDPMixin(object): # 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: - m = f""" - Amp level {self.amp_level} with DataParallel is not supported. - See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227. - We recommend you switch to ddp if you want to use amp - """ - raise MisconfigurationException(m) + if self.amp_level == 'O2': + m = f""" + Amp level {self.amp_level} with DataParallel is not supported. + See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227. + We recommend you switch to ddp if you want to use amp + """ + raise MisconfigurationException(m) + else: + model, optimizers = model.configure_apex(amp, model, self.optimizers, self.amp_level) # create list of device ids device_ids = self.data_parallel_device_ids