From 56fa2075a528a7a8744f4f803df902f02ef36372 Mon Sep 17 00:00:00 2001
From: David Kossnick <david@angle-tech.com>
Date: Wed, 23 Oct 2019 03:11:18 -0700
Subject: [PATCH] Move `global_step` incrementing (#412)

* Move global_step incrementing to the end of a batch loop, per https://github.com/williamFalcon/pytorch-lightning/issues/411

* Move met_batch_limit condition to the end

* cleanup whitespace

* Update train_loop_mixin.py
---
 pytorch_lightning/trainer/train_loop_mixin.py | 20 +++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/pytorch_lightning/trainer/train_loop_mixin.py b/pytorch_lightning/trainer/train_loop_mixin.py
index e51a43d9..1dad28ae 100644
--- a/pytorch_lightning/trainer/train_loop_mixin.py
+++ b/pytorch_lightning/trainer/train_loop_mixin.py
@@ -74,23 +74,17 @@ class TrainerTrainLoopMixin(object):
         # run epoch
         for batch_nb, batch in enumerate(self.get_train_dataloader()):
             self.batch_nb = batch_nb
-            self.global_step += 1
 
             model = self.get_model()
             model.global_step = self.global_step
 
-            # stop when the flag is changed or we've gone past the amount
-            #  requested in the batches
-            self.total_batch_nb += 1
-            met_batch_limit = batch_nb >= self.nb_training_batches
-            if met_batch_limit:
-                break
-
             # ---------------
             # RUN TRAIN STEP
             # ---------------
             output = self.run_training_batch(batch, batch_nb)
             batch_result, grad_norm_dic, batch_step_metrics = output
+
+            # when returning -1 from train_step, we end epoch early
             early_stop_epoch = batch_result == -1
 
             # ---------------
@@ -116,10 +110,20 @@ class TrainerTrainLoopMixin(object):
                 # logs user requested information to logger
                 self.log_metrics(batch_step_metrics, grad_norm_dic)
 
+            self.global_step += 1
+            self.total_batch_nb += 1
+
             # end epoch early
+            # stop when the flag is changed or we've gone past the amount
+            # requested in the batches
             if early_stop_epoch or self.fast_dev_run:
                 break
 
+            # stop epoch if we limited nb batches
+            met_batch_limit = batch_nb >= self.nb_training_batches
+            if met_batch_limit:
+                break
+
         # epoch end hook
         if self.is_function_implemented('on_epoch_end'):
             model = self.get_model()