Move model to cuda before creating optimizer (#554)

This commit is contained in:
Anton Bakhtin 2019-11-26 22:35:38 -05:00 committed by William Falcon
parent f2191b0cdf
commit fea7cc87f6
1 changed files with 2 additions and 2 deletions

View File

@ -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)