From e350a7db076e32633904fc769e4ae79c46bbaf7d Mon Sep 17 00:00:00 2001 From: William Falcon Date: Mon, 11 Nov 2019 22:58:11 -0500 Subject: [PATCH] Enable apex O2 + dp (#493) * remove O2 crash * remove O2 crash * bananas --- pytorch_lightning/trainer/dp_mixin.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) 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