Move model to cuda before creating optimizer (#554)

This commit is contained in:
Anton Bakhtin authored and William Falcon committed 2019-11-26 22:35:38 -05:00
1 parent f2191b0cdf
commit fea7cc87f6
1 file changed
+2 -2
+2 -2
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)