diff --git a/pytorch_lightning/trainer/train_loop_mixin.py b/pytorch_lightning/trainer/train_loop_mixin.py index 68c34cf4..7e9fde01 100644 --- a/pytorch_lightning/trainer/train_loop_mixin.py +++ b/pytorch_lightning/trainer/train_loop_mixin.py @@ -155,7 +155,7 @@ class TrainerTrainLoopMixin(object): all_log_metrics = [] if batch is None: - return 0, grad_norm_dic + return 0, grad_norm_dic, {} # hook if self.is_function_implemented('on_batch_start'): @@ -163,7 +163,7 @@ class TrainerTrainLoopMixin(object): response = model_ref.on_batch_start(batch) if response == -1: - return -1, grad_norm_dic + return -1, grad_norm_dic, {} splits = [batch] if self.truncated_bptt_steps is not None: