From 1b273a32eea075295b0ea73b210cb43017d22194 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Wed, 24 Jul 2019 14:11:05 -0400 Subject: [PATCH] fixed amp bug --- pytorch_lightning/models/trainer.py | 8 ++++++++ tests/debug.py | 2 +- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index 68ffe8c5..b2f9b66e 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -460,6 +460,14 @@ class Trainer(TrainerIO): model.cuda(self.data_parallel_device_ids[0]) + # 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': + 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) + # run through amp wrapper if self.use_amp: diff --git a/tests/debug.py b/tests/debug.py index 876075c5..7c8543d3 100644 --- a/tests/debug.py +++ b/tests/debug.py @@ -111,7 +111,7 @@ def main(): max_nb_epochs=1, gpus=[0, 1], distributed_backend='dp', - use_amp=False + use_amp=True ) result = trainer.fit(model)