diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index 7b763136..f82d4ca2 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -399,7 +399,7 @@ class Trainer(TrainerIO):
         :return:
         """
         ref_model = model
-        if self.on_gpu:
+        if self.data_parallel:
             ref_model = model.module
 
         ref_model.trainer = self