From 64de447545ba7d17da61c321c4d182bb2d2339dc Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Fri, 26 Jul 2019 22:07:02 -0400
Subject: [PATCH] fixed hpc save, load. cleaned apu

---
 pytorch_lightning/root_module/root_module.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py
index a25f7f6f..b49dd3af 100644
--- a/pytorch_lightning/root_module/root_module.py
+++ b/pytorch_lightning/root_module/root_module.py
@@ -108,13 +108,12 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         else:
             checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)
 
+        # load the state_dict on the model automatically
         model = cls(hparams)
+        model.load_state_dict(checkpoint['state_dict'])
 
         # give model a chance to load something
-        model.load_model_specific(checkpoint)
-
-        # load the state_dict on the model automatically
-        model.load_state_dict(checkpoint['state_dict'])
+        model.on_load_checkpoint(checkpoint)
 
         return model