Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

refactored model tests

This commit is contained in:
William Falcon 2019-07-24 13:56:49 -04:00
parent ef843d5f96
commit ecb68b52f8
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
pytorch_lightning/models/trainer.py
View File

@ -458,9 +458,10 @@ class Trainer(TrainerIO):
# filter out the weights that were done on gpu so we can load on good old cpus # filter out the weights that were done on gpu so we can load on good old cpus
self.optimizers = model.configure_optimizers() self.optimizers = model.configure_optimizers()
model.cuda(self.data_parallel_device_ids[0])
# run through amp wrapper # run through amp wrapper
if self.use_amp: if self.use_amp:
model.cuda(self.data_parallel_device_ids[0])
# An example # An example
model, optimizers = amp.initialize( model, optimizers = amp.initialize(