Move model to cuda before creating optimizer (#554)
This commit is contained in:
parent
f2191b0cdf
commit
fea7cc87f6
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue