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

modified single gpu init

This commit is contained in:
William Falcon 2019-07-14 17:01:18 -04:00
parent e520297781
commit 849f52b7a6
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
pytorch_lightning/models/trainer.py
View File

@ -320,7 +320,7 @@ class Trainer(TrainerIO):
def single_gpu_train(self, model): def single_gpu_train(self, model):
torch.cuda.set_device(0) torch.cuda.set_device(0)
model = model.cuda(0) model.cuda(0)
# CHOOSE OPTIMIZER # CHOOSE OPTIMIZER
# 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