From ffa7a0dbab42df2e0dcadde1ee5b4a519ae80481 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Fri, 26 Jul 2019 22:26:55 -0400
Subject: [PATCH] added clean slurm save load test

---
 pytorch_lightning/root_module/model_saving.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py
index 142d2b337b..361a08548b 100644
--- a/pytorch_lightning/root_module/model_saving.py
+++ b/pytorch_lightning/root_module/model_saving.py
@@ -181,6 +181,7 @@ class TrainerIO(object):
         # call model hook
         model.on_hpc_load(checkpoint)
 
+        self.model = model
 
     def max_ckpt_in_folder(self, path):
         files = os.listdir(path)