From 24a3246bc195e6c3b7c057e27b83ba636d8f252f Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Thu, 25 Jul 2019 10:56:03 -0400
Subject: [PATCH] updated test models with lazy decorators

---
 pytorch_lightning/__init__.py                 |  3 +-
 .../lightning_module_template.py              | 31 +++++--------------
 .../sample_model_template/model_template.py   |  4 +--
 pytorch_lightning/root_module/__init__.py     |  1 +
 pytorch_lightning/root_module/root_module.py  |  8 ++---
 5 files changed, 16 insertions(+), 31 deletions(-)

diff --git a/pytorch_lightning/__init__.py b/pytorch_lightning/__init__.py
index 5893eb4a..e7a03ba8 100644
--- a/pytorch_lightning/__init__.py
+++ b/pytorch_lightning/__init__.py
@@ -1,2 +1,3 @@
 from .models import Trainer
-from .root_module.root_module import LightningModule
\ No newline at end of file
+from .root_module.root_module import LightningModule
+from .root_module.decorators import data_loader
\ No newline at end of file
diff --git a/pytorch_lightning/examples/new_project_templates/lightning_module_template.py b/pytorch_lightning/examples/new_project_templates/lightning_module_template.py
index c6b5ccd1..5a6995ec 100644
--- a/pytorch_lightning/examples/new_project_templates/lightning_module_template.py
+++ b/pytorch_lightning/examples/new_project_templates/lightning_module_template.py
@@ -10,6 +10,7 @@ from torch import optim
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
 
+import pytorch_lightning as ptl
 from pytorch_lightning.root_module.root_module import LightningModule
 
 
@@ -200,35 +201,17 @@ class LightningTemplateModel(LightningModule):
 
         return loader
 
-    @property
+    @ptl.data_loader
     def tng_dataloader(self):
-        if self._tng_dataloader is None:
-            try:
-                self._tng_dataloader = self.__dataloader(train=True)
-            except Exception as e:
-                print(e)
-                raise e
-        return self._tng_dataloader
+        return self.__dataloader(train=True)
 
-    @property
+    @ptl.data_loader
     def val_dataloader(self):
-        if self._val_dataloader is None:
-            try:
-                self._val_dataloader = self.__dataloader(train=False)
-            except Exception as e:
-                print(e)
-                raise e
-        return self._val_dataloader
+        return self.__dataloader(train=False)
 
-    @property
+    @ptl.data_loader
     def test_dataloader(self):
-        if self._test_dataloader is None:
-            try:
-                self._test_dataloader = self.__dataloader(train=False)
-            except Exception as e:
-                print(e)
-                raise e
-        return self._test_dataloader
+        return self.__dataloader(train=False)
 
     @staticmethod
     def add_model_specific_args(parent_parser, root_dir): # pragma: no cover
diff --git a/pytorch_lightning/models/sample_model_template/model_template.py b/pytorch_lightning/models/sample_model_template/model_template.py
index 44c57570..10f12c59 100644
--- a/pytorch_lightning/models/sample_model_template/model_template.py
+++ b/pytorch_lightning/models/sample_model_template/model_template.py
@@ -1,6 +1,6 @@
 import torch.nn as nn
 import numpy as np
-from pytorch_lightning.root_module.root_module import LightningModule
+from pytorch_lightning import LightningModule
 from test_tube import HyperOptArgumentParser
 from torchvision.datasets import MNIST
 import torchvision.transforms as transforms
@@ -149,7 +149,7 @@ class ExampleModel1(LightningModule):
 
         return loader
 
-    @property
+    @data_loader
     def tng_dataloader(self):
         if self._tng_dataloader is None:
             try:
diff --git a/pytorch_lightning/root_module/__init__.py b/pytorch_lightning/root_module/__init__.py
index e69de29b..63ea400e 100644
--- a/pytorch_lightning/root_module/__init__.py
+++ b/pytorch_lightning/root_module/__init__.py
@@ -0,0 +1 @@
+from .decorators import data_loader
\ No newline at end of file
diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py
index d6e0ec96..0cb37c23 100644
--- a/pytorch_lightning/root_module/root_module.py
+++ b/pytorch_lightning/root_module/root_module.py
@@ -3,7 +3,7 @@ from pytorch_lightning.root_module.memory import ModelSummary
 from pytorch_lightning.root_module.grads import GradInformation
 from pytorch_lightning.root_module.model_saving import ModelIO, load_hparams_from_tags_csv
 from pytorch_lightning.root_module.hooks import ModelHooks
-from pytorch_lightning.root_module.decorators import data_loader
+import pytorch_lightning as ptl
 
 
 class LightningModule(GradInformation, ModelIO, ModelHooks):
@@ -84,7 +84,7 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         for param in self.parameters():
             param.requires_grad = True
 
-    @data_loader
+    @ptl.data_loader
     def tng_dataloader(self):
         """
         Implement a function to load an h5py of this data
@@ -92,7 +92,7 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         """
         raise NotImplementedError
 
-    @data_loader
+    @ptl.data_loader
     def test_dataloader(self):
         """
         Implement a function to load an h5py of this data
@@ -100,7 +100,7 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         """
         raise NotImplementedError
 
-    @data_loader
+    @ptl.data_loader
     def val_dataloader(self):
         """
         Implement a function to load an h5py of this data