pointer to trainer in model

pytorch_lightning/models/trainer.py

@@ -86,7 +86,7 @@ class Trainer(TrainerIO):
             self.test_percent_check = overfit_pct
 
     def __is_function_implemented(self, f_name):
-        f_op = getattr(self, f_name, None)
+        f_op = getattr(self.model, f_name, None)
         return callable(f_op)
 
     @property
@@ -195,6 +195,7 @@ class Trainer(TrainerIO):
     # -----------------------------
     def fit(self, model):
         self.model = model
+        model.trainer = self
 
         # transfer data loaders from model
         self.__get_dataloaders(model)
@@ -266,10 +267,6 @@ class Trainer(TrainerIO):
                 if met_batch_limit:
                     break
 
-                # give model a chance to end epoch early
-                if self.model.should_stop_epoch(data_batch):
-                    break
-
                 # ---------------
                 # RUN TRAIN STEP
                 # ---------------
@@ -331,7 +328,9 @@ class Trainer(TrainerIO):
 
         # hook
         if self.__is_function_implemented('on_batch_start'):
-            self.model.on_batch_start()
+            response = self.model.on_batch_start(data_batch)
+            if response == -1:
+                return
 
         if self.enable_tqdm:
             self.prog_bar.update(1)

pytorch_lightning/root_module/hooks.py

@@ -1,7 +1,7 @@
 import torch
 
 class ModelHooks(torch.nn.Module):
-    def on_batch_start(self):
+    def on_batch_start(self, data_batch):
         pass
 
     def on_batch_end(self):
@@ -19,5 +19,3 @@ class ModelHooks(torch.nn.Module):
     def on_post_performance_check(self):
         pass
 
-    def should_stop_epoch(self, data_batch):
-        return False

pytorch_lightning/root_module/root_module.py

@@ -24,6 +24,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         self.overfit = hparams.overfit
         self.gradient_clip = hparams.gradient_clip
         self.num = 2
+        self.trainer = None
 
         # track if gpu was requested for checkpointing
         self.on_gpu = False

setup.py

@@ -7,7 +7,7 @@ from setuptools import setup, find_packages
 # http://blog.ionelmc.ro/2014/05/25/python-packaging/
 setup(
     name="pytorch-lightning",
-    version='0.1.dev12',
+    version='0.1.dev14',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",