From 31d9062b3a92d0cf545f43493e1ad1733076e12f Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Mon, 8 Jul 2019 18:55:05 -0400
Subject: [PATCH] moved sampler

---
 pytorch_lightning/models/trainer.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index bb63e35b..3daa759d 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -283,7 +283,6 @@ class Trainer(TrainerIO):
                 model, self.optimizers, opt_level=self.amp_level,
             )
             self.optimizers = optimizers
-            model.trainer = self
 
         # when using gpus, first thing we do is spawn a new process between each worker
         # applies to single gpu, multi-gpu and multi-nodes
@@ -385,6 +384,8 @@ class Trainer(TrainerIO):
         if self.on_gpu:
             ref_model = model.module
 
+        ref_model.trainer = self
+
         # set local properties on the model
         ref_model.on_gpu = self.on_gpu