modified single gpu init

2019-07-14 17:01:18 -04:00 · 2019-07-14 17:01:18 -04:00 · 849f52b7a6
parent e520297781
commit 849f52b7a6
1 changed files with 1 additions and 1 deletions
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@ -320,7 +320,7 @@ class Trainer(TrainerIO):

    def single_gpu_train(self, model):
        torch.cuda.set_device(0)
-        model = model.cuda(0)
+        model.cuda(0)

        # CHOOSE OPTIMIZER
        # filter out the weights that were done on gpu so we can load on good old cpus