From 6876f60098c92f4c55a20bb4022fe3383e102cd8 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Sat, 13 Jul 2019 22:21:17 -0400
Subject: [PATCH] merge

---
 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 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