From fea7cc87f6d961157435559d46ebcabf7edf5dab Mon Sep 17 00:00:00 2001 From: Anton Bakhtin <7470747+akhti@users.noreply.github.com> Date: Tue, 26 Nov 2019 22:35:38 -0500 Subject: [PATCH] Move model to cuda before creating optimizer (#554) --- pytorch_lightning/trainer/dp_mixin.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pytorch_lightning/trainer/dp_mixin.py b/pytorch_lightning/trainer/dp_mixin.py index 6bc08704..efa10b49 100644 --- a/pytorch_lightning/trainer/dp_mixin.py +++ b/pytorch_lightning/trainer/dp_mixin.py @@ -63,12 +63,12 @@ class TrainerDPMixin(object): return batch def single_gpu_train(self, model): + model.cuda(self.root_gpu) + # CHOOSE OPTIMIZER # allow for lr schedulers as well self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers()) - model.cuda(self.root_gpu) - if self.use_amp: # An example model, optimizers = model.configure_apex(amp, model, self.optimizers, self.amp_level)