From df4ac681ed0faa5d349b28dd4963b73e5a4a04d6 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 26 Jun 2019 18:08:46 -0400
Subject: [PATCH] removed self.model refs

---
 pytorch_lightning/utils/pt_callbacks.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/pytorch_lightning/utils/pt_callbacks.py b/pytorch_lightning/utils/pt_callbacks.py
index 6c388813..46e6195d 100644
--- a/pytorch_lightning/utils/pt_callbacks.py
+++ b/pytorch_lightning/utils/pt_callbacks.py
@@ -1,5 +1,6 @@
 import numpy as np
 import os, shutil
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDataParallel
 
 
 class Callback(object):
@@ -33,6 +34,8 @@ class Callback(object):
         self.params = params
 
     def set_model(self, model):
+        if type(model) is LightningDataParallel:
+            model = model.module
         self.model = model
 
     def on_epoch_begin(self, epoch, logs=None):