From 849f52b7a6005e10a824ab9c73b613bdd0eb3883 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Sun, 14 Jul 2019 17:01:18 -0400 Subject: [PATCH] modified single gpu init --- pytorch_lightning/models/trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index 52dcb4be26..479456ccea 100644 --- 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