Repository: Lightning-AI/pytorch-lightning
Pull request: Fix inconsistent outputs in `on_*_end` and `*_end`
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
 Please also include relevant motivation and context.
 List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->
Fixes #6921

Several changes:
- Move `training_epoch_end` (which actual handles model epoch end) from logger connector to `training_loop.py`
- Use one static method (`_prepare_outputs`) in `training_loop.py` to unpack all `Result` object outputs from both `epoch_end` and `batch_end`
- Pass these unpacked outputs to __both__ `on_train_epoch_end` and `training_epoch_end` hooks (previously, `training_epoch_end` behaved differently)
- Moved evaluation output handling to `trainer.py` so that both `evaluation_epoch_end` and `on_evaluation_epoch_end` get the same `outputs` argument
- Changed order of calls to `training_epoch_end` and `on_train_epoch_end` to be consistent with the docs

Previously, running the code from the issue, the user saw this output:

```python
callback on_validation_batch_end: {'loss': tensor(1.9916), 'foo': 'from_val_step'}
validation_epoch_end: [{'loss': tensor(1.9916), 'foo': 'from_val_step'}]
callback on_validation_epoch_end: [[{'loss': tensor(1.9916), 'foo': 'from_val_step'}]]
on_validation_epoch_end: [[{'loss': tensor(1.9916), 'foo': 'from_val_step'}]]
callback on_train_batch_end: [[{'extra': {'foo': 123}, 'minimize': tensor(0.2295)}]]
callback on_train_epoch_end: [[[{'extra': {'foo': 123}, 'minimize': tensor(0.2295)}]]]
on_train_epoch_end: [[[{'extra': {'foo': 123}, 'minimize': tensor(0.2295)}]]]
training_epoch_end: [{'foo': 123, 'loss': tensor(0.2295)}]
callback on_validation_batch_end: {'loss': tensor(2.2656), 'foo': 'from_val_step'}
validation_epoch_end: [{'loss': tensor(2.2656), 'foo': 'from_val_step'}]
callback on_validation_epoch_end: [[{'loss': tensor(2.2656), 'foo': 'from_val_step'}]]
on_validation_epoch_end: [[{'loss': tensor(2.2656), 'foo': 'from_val_step'}]]
```

Now sees this:
```python
callback on_validation_batch_end: {'loss': tensor(0.6422), 'foo': 'from_val_step'}
validation_epoch_end: [{'loss': tensor(0.6422), 'foo': 'from_val_step'}]
callback on_validation_epoch_end: [{'loss': tensor(0.6422), 'foo': 'from_val_step'}]
on_validation_epoch_end: [{'loss': tensor(0.6422), 'foo': 'from_val_step'}]
callback on_train_batch_end: {'foo': 123, 'loss': tensor(1.2562)}
training_epoch_end: [{'foo': 123, 'loss': tensor(1.2562)}]
callback on_train_epoch_end: [{'foo': 123, 'loss': tensor(1.2562)}]
on_train_epoch_end: [{'foo': 123, 'loss': tensor(1.2562)}]
callback on_validation_batch_end: {'loss': tensor(0.4613), 'foo': 'from_val_step'}
validation_epoch_end: [{'loss': tensor(0.4613), 'foo': 'from_val_step'}]
callback on_validation_epoch_end: [{'loss': tensor(0.4613), 'foo': 'from_val_step'}]
on_validation_epoch_end: [{'loss': tensor(0.4613), 'foo': 'from_val_step'}]
```

## Before submitting
- [x] Was this discussed/approved via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your PR does only one thing, instead of bundling different changes together?
- [x] Did you make sure to update the documentation with your changes? (if necessary)
- [x] Did you write any new necessary tests? (not for typos and docs)
- [x] Did you verify new and existing tests pass locally with your changes?
- [x] Did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- For CHANGELOG separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review
Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

 - [x] Is this pull request ready for review? (if not, please submit in draft mode)
 - [x] Check that all items from **Before submitting** are resolved
 - [x] Make sure the title is self-explanatory and the description concisely explains the PR
 - [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?
Make sure you had fun coding 🙃


Problem statement:
Inconsistent `outputs` format between `training_epoch_end` and `on_train_epoch_end`
## 🐛 Bug

The outputs object for `on_train_epoch_end` should not include the `extra` field

### To Reproduce

```python
def test_bug(tmpdir):
    class TestModel(BoringModel):
        def training_step(self, batch, batch_idx):
            output = self(batch)
            loss = self.loss(batch, output)
            return {"loss": loss, "foo": 123}

        def training_epoch_end(self, outputs):
            print("training_epoch_end:", outputs)

        def on_train_epoch_end(self, outputs):
            print("on_train_epoch_end:", outputs)

    class TestCallback(Callback):
        def on_train_epoch_end(self, trainer, pl_module, outputs):
            print("callback on_train_epoch_end:", outputs)

    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, callbacks=[TestCallback()], progress_bar_refresh_rate=0)
    trainer.fit(TestModel())
```

```console
callback on_train_epoch_end: [[[{'extra': {'foo': 123}, 'minimize': tensor(1.1792)}]]]
on_train_epoch_end: [[[{'extra': {'foo': 123}, 'minimize': tensor(1.1792)}]]]
training_epoch_end: [{'foo': 123, 'loss': tensor(1.1792)}]
```

### Expected behavior

`on_train_epoch_end` acts as `training_epoch_end`

### Environment

master

### Additional context

Reported by Marek O in Slack

Hints from issue discussion:
Hi, I'm Marek from slack, thank you for the reproducible! 

I'll add extended test that shows also difference in `on_xxx_batch_end()` between train and val/test: 
### Test
```
import os
from typing import Any

import torch
from pl_examples.bug_report_model import BoringModel, RandomDataset
from pytorch_lightning import Trainer, Callback, LightningModule


def test_bug(tmpdir):
    class TestModel(BoringModel):
        def training_step(self, batch, batch_idx):
            output = self(batch)
            loss = self.loss(batch, output)
            return {"loss": loss, "foo": 123}

        def validation_step(self, batch, batch_idx):
            output = self(batch)
            loss = self.loss(batch, output)
            return {"loss": loss, "foo": 'from_val_step'}

        def training_epoch_end(self, outputs):
            print("training_epoch_end:", outputs)

        def validation_epoch_end(self, outputs) -> None:
            print("validation_epoch_end:", outputs)

        def on_train_epoch_end(self, outputs):
            print("on_train_epoch_end:", outputs)

        def on_validation_epoch_end(self) -> None:  #FIXME should have (optional) outputs as validation_epoch_end(self, outputs)
            #print("on_validation_epoch_end:", outputs)
            print("on_validation_epoch_end:", 'FIXME should have outputs')

    class TestCallback(Callback):
        def on_train_epoch_end(self, trainer, pl_module, outputs):
            print("callback on_train_epoch_end:", outputs)

        def on_train_batch_end(self, trainer, pl_module: LightningModule, outputs: Any, batch: Any, batch_idx: int, dataloader_idx: int) -> None:
            print('callback on_train_batch_end:', outputs)


        def on_validation_batch_end(self, trainer, pl_module: LightningModule, outputs: Any, batch: Any, batch_idx: int, dataloader_idx: int) -> None:
            print("callback on_validation_batch_end:", outputs)

    # fake data
    train_data = torch.utils.data.DataLoader(RandomDataset(32, 64))
    val_data = torch.utils.data.DataLoader(RandomDataset(32, 64))
    test_data = torch.utils.data.DataLoader(RandomDataset(32, 64))

    # model
    model = TestModel()
    trainer = Trainer(
        default_root_dir=tmpdir,
        fast_dev_run=False,
        limit_train_batches=1,
        limit_val_batches=1,
        max_epochs=1,
        callbacks=[TestCallback()],
        progress_bar_refresh_rate=0,
        weights_summary=None,
    )

    trainer.fit(model, train_data, val_data)
    trainer.test(test_dataloaders=test_data)



if __name__ == '__main__':
    test_bug('/tmp')
```

### output: 
```
callback on_validation_batch_end: {'loss': tensor(2.7839), 'foo': 'from_val_step'}
validation_epoch_end: [{'loss': tensor(2.7839), 'foo': 'from_val_step'}]
on_validation_epoch_end: FIXME should have outputs
callback on_train_batch_end: [[{'extra': {'foo': 123}, 'minimize': tensor(1.2255)}]]
callback on_train_epoch_end: [[[{'extra': {'foo': 123}, 'minimize': tensor(1.2255)}]]]
on_train_epoch_end: [[[{'extra': {'foo': 123}, 'minimize': tensor(1.2255)}]]]
training_epoch_end: [{'foo': 123, 'loss': tensor(1.2255)}]
callback on_validation_batch_end: {'loss': tensor(2.1134), 'foo': 'from_val_step'}
validation_epoch_end: [{'loss': tensor(2.1134), 'foo': 'from_val_step'}]
on_validation_epoch_end: FIXME should have outputs
```
### Expected behavior: 
`on_validation_epoch_end(self)` (and on_test_epoch_end(self)) have same API as `on_train_epoch_end(self, outputs)`

@carmocca i'm running into this bug as well in #6944 
The issue is that the pseudocode here is not right! https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#hooks

```
# end training epoch
outs = training_epoch_end(outs)
on_train_epoch_end(outs)
on_epoch_end()
```
From the code: https://github.com/PyTorchLightning/pytorch-lightning/blob/b85cfbe8f350a89c93ff967fa416859be7ebb4f3/pytorch_lightning/trainer/training_loop.py#L480-L484

the model's `training_epoch_end` is called inside of the logger connector: https://github.com/PyTorchLightning/pytorch-lightning/blob/b85cfbe8f350a89c93ff967fa416859be7ebb4f3/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py#L331-L339

which happens after `on_train_epoch_end` is called

@edenafek @tchaton 

Thanks for looking @ananthsub! So we should fix the order and maybe also call `logger_connector.training_epoch_end` in the training loop instead so the logic is a tiny bit cleaner.

Commit message: Update CHANGELOG.md

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0f96664db72b0..6794286919ff9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -240,6 +240,15 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed `AttributeError for `require_backward_grad_sync` when running manual optimization with sharded plugin ([#6915](https://github.com/PyTorchLightning/pytorch-lightning/pull/6915))
 
 
+- Fixed a bug where `LightningModule.train_epoch_end` was called after the `on_train_end_epoch` hook ([#6969](https://github.com/PyTorchLightning/pytorch-lightning/pull/6969))
+
+
+- Fixed a bug where the outputs object passed to `LightningModule.train_epoch_end` was different from the object passed to the `on_train_end_epoch` hook ([#6969](https://github.com/PyTorchLightning/pytorch-lightning/pull/6969))
+
+
+- Fixed a bug where the outputs passed to `train_batch_end` would be lists even when using a single optimiser and no truncated backprop through time steps ([#6969](https://github.com/PyTorchLightning/pytorch-lightning/pull/6969))
+
+
 ## [1.2.7] - 2021-04-06
 
 ### Fixed
diff --git a/docs/source/common/lightning_module.rst b/docs/source/common/lightning_module.rst
index 95f7c0a6dcb7e..912c5e4d30944 100644
--- a/docs/source/common/lightning_module.rst
+++ b/docs/source/common/lightning_module.rst
@@ -1046,7 +1046,7 @@ This is the pseudocode to describe how all the hooks are called during a call to
                 val_loop()
 
         # end training epoch
-        outs = training_epoch_end(outs)
+        training_epoch_end(outs)
         on_train_epoch_end(outs)
         on_epoch_end()
 
diff --git a/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py b/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py
index 8fd716ee64ad7..5ec4a70099145 100644
--- a/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py
+++ b/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py
@@ -26,8 +26,6 @@
 from pytorch_lightning.trainer.connectors.logger_connector.metrics_holder import MetricsHolder
 from pytorch_lightning.trainer.states import RunningStage, TrainerState
 from pytorch_lightning.utilities import DeviceType, flatten_dict
-from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.model_helpers import is_overridden
 
 
 class LoggerConnector:
@@ -328,16 +326,11 @@ def on_train_epoch_end(self):
         # inform cached logger connector epoch finished
         self.cached_results.has_batch_loop_finished = True
 
-    def log_train_epoch_end_metrics(self, epoch_output, num_optimizers):
+    def log_train_epoch_end_metrics(self, epoch_output):
         # epoch output is a list. Each item in that list has all the outputs per optimizer
         # epoch_output[optimizer_idx][training_step_idx][tbptt_index]
         # remember that not using truncated backprop is equivalent with truncated back prop of len(1)
 
-        model = self.trainer.lightning_module
-
-        # lightning module hook
-        self.training_epoch_end(model, epoch_output, num_optimizers)
-
         # log/aggregate metrics automatically
         epoch_log_metrics, epoch_progress_bar_metrics = self.__auto_reduce_results_on_epoch_end(epoch_output)
 
@@ -365,29 +358,6 @@ def log_train_epoch_end_metrics(self, epoch_output, num_optimizers):
         # reset epoch loop result for next epoch
         self.cached_results.reset()
 
-    def training_epoch_end(self, model, epoch_output, num_optimizers):
-        if not is_overridden('training_epoch_end', model=model):
-            return
-
-        # run training_epoch_end
-        # refresh the result for custom logging at the epoch level
-        model._current_fx_name = 'training_epoch_end'
-        epoch_output = self.__prepare_epoch_end_inputs(epoch_output)
-
-        if num_optimizers == 1 or not self.trainer.train_loop.automatic_optimization:
-            epoch_output = epoch_output[0]
-
-        # lightningmodule hook
-        epoch_output = model.training_epoch_end(epoch_output)
-
-        if epoch_output is not None:
-            raise MisconfigurationException(
-                'training_epoch_end expects a return of None. '
-                'HINT: remove the return statement in training_epoch_end'
-            )
-        # capture logging
-        self.trainer.logger_connector.cache_logged_metrics()
-
     def __auto_reduce_results_on_epoch_end(self, epoch_output):
         epoch_log_metrics = {}
         epoch_progress_bar_metrics = {}
@@ -413,33 +383,6 @@ def __auto_reduce_results_on_epoch_end(self, epoch_output):
 
         return epoch_log_metrics, epoch_progress_bar_metrics
 
-    def __prepare_epoch_end_inputs(self, epoch_output):
-        """
-        Pulls out only the "extra" information for epoch end
-
-        Return:
-            a single list, each element per optimizer then batch then time
-        """
-        gathered_epoch_outputs = []
-        for opt_outputs in epoch_output:
-            # gather across time first
-            time_gathered_outputs = []
-            for tbptt_outs in opt_outputs:
-                result = []
-                for x in tbptt_outs:
-                    out = x.extra
-                    out['loss'] = x.minimize
-                    result.append(out)
-
-                # when time = 0, pass in the literal dict instead of array
-                if len(result) == 1:
-                    result = result[0]
-                time_gathered_outputs.append(result)
-
-            gathered_epoch_outputs.append(time_gathered_outputs)
-
-        return gathered_epoch_outputs
-
     def log_train_step_metrics(self, batch_output):
         if self.trainer.train_loop.should_accumulate() and self.trainer.train_loop.automatic_optimization:
             return
diff --git a/pytorch_lightning/trainer/evaluation_loop.py b/pytorch_lightning/trainer/evaluation_loop.py
index 8b7543e6bf50f..36e51645c3842 100644
--- a/pytorch_lightning/trainer/evaluation_loop.py
+++ b/pytorch_lightning/trainer/evaluation_loop.py
@@ -186,12 +186,12 @@ def evaluation_step_end(self, *args, **kwargs):
             output = self.trainer.call_hook('validation_step_end', *args, **kwargs)
         return output
 
-    def evaluation_epoch_end(self):
+    def evaluation_epoch_end(self, outputs):
         # unset dataloder_idx in model
         self.trainer.logger_connector.evaluation_epoch_end()
 
         # call the model epoch end
-        deprecated_results = self.__run_eval_epoch_end(self.num_dataloaders)
+        deprecated_results = self.__run_eval_epoch_end(outputs)
 
         # enable returning anything
         for i, r in enumerate(deprecated_results):
@@ -205,46 +205,40 @@ def log_epoch_metrics_on_evaluation_end(self):
         eval_loop_results = self.trainer.logger_connector.get_evaluate_epoch_results()
         return eval_loop_results
 
-    def __run_eval_epoch_end(self, num_dataloaders):
+    def __run_eval_epoch_end(self, outputs):
         model = self.trainer.lightning_module
 
-        # with a single dataloader don't pass an array
-        outputs = self.outputs
-
-        eval_results = outputs
-        if num_dataloaders == 1:
-            eval_results = outputs[0]
-
+        epoch_end_output = None
         user_reduced = False
 
         if self.trainer.testing:
             if is_overridden('test_epoch_end', model=model):
                 model._current_fx_name = 'test_epoch_end'
-                eval_results = model.test_epoch_end(eval_results)
+                epoch_end_output = model.test_epoch_end(outputs)
                 user_reduced = True
 
         else:
             if is_overridden('validation_epoch_end', model=model):
                 model._current_fx_name = 'validation_epoch_end'
-                eval_results = model.validation_epoch_end(eval_results)
+                epoch_end_output = model.validation_epoch_end(outputs)
                 user_reduced = True
 
         # capture logging
         self.trainer.logger_connector.cache_logged_metrics()
         # depre warning
-        if eval_results is not None and user_reduced:
+        if epoch_end_output is not None and user_reduced:
             step = 'testing_epoch_end' if self.trainer.testing else 'validation_epoch_end'
             self.warning_cache.warn(
                 f'The {step} should not return anything as of 9.1.'
                 ' To log, use self.log(...) or self.write(...) directly in the LightningModule'
             )
 
-        if not isinstance(eval_results, list):
-            eval_results = [eval_results]
+        if not isinstance(outputs, list):
+            outputs = [outputs]
 
-        self.trainer.logger_connector._track_callback_metrics(eval_results)
+        self.trainer.logger_connector._track_callback_metrics(outputs)
 
-        return eval_results
+        return outputs
 
     def __gather_epoch_end_eval_results(self, outputs):
         eval_results = []
@@ -307,18 +301,13 @@ def store_predictions(self, output, batch_idx, dataloader_idx):
         # track debug metrics
         self.trainer.dev_debugger.track_eval_loss_history(batch_idx, dataloader_idx, output)
 
-    def on_evaluation_epoch_end(self, *args, **kwargs):
+    def on_evaluation_epoch_end(self, outputs):
         # call the callback hook
-        self.call_on_evaluation_epoch_end_hook()
+        self.call_on_evaluation_epoch_end_hook(outputs)
 
         self.trainer.call_hook('on_epoch_end')
 
-    def call_on_evaluation_epoch_end_hook(self):
-        outputs = self.outputs
-
-        # free memory
-        self.outputs = []
-
+    def call_on_evaluation_epoch_end_hook(self, outputs):
         model_ref = self.trainer.lightning_module
         hook_name = "on_test_epoch_end" if self.trainer.testing else "on_validation_epoch_end"
 
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
index e6bf36df92a01..ede64154e1ed7 100644
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -711,11 +711,20 @@ def run_evaluation(self, on_epoch=False):
             # store batch level output per dataloader
             self.evaluation_loop.outputs.append(dl_outputs)
 
+        outputs = self.evaluation_loop.outputs
+
+        # reset outputs
+        self.evaluation_loop.outputs = []
+
+        # with a single dataloader don't pass an array
+        if self.evaluation_loop.num_dataloaders == 1:
+            outputs = outputs[0]
+
         # lightning module method
-        deprecated_eval_results = self.evaluation_loop.evaluation_epoch_end()
+        deprecated_eval_results = self.evaluation_loop.evaluation_epoch_end(outputs)
 
         # hook
-        self.evaluation_loop.on_evaluation_epoch_end()
+        self.evaluation_loop.on_evaluation_epoch_end(outputs)
 
         # update epoch-level lr_schedulers
         if on_epoch:
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
index d2749733812b3..7c77363173f09 100644
--- a/pytorch_lightning/trainer/training_loop.py
+++ b/pytorch_lightning/trainer/training_loop.py
@@ -185,8 +185,12 @@ def on_train_epoch_start(self, epoch):
         self.trainer.call_hook("on_train_epoch_start")
 
     def on_train_batch_end(self, epoch_output, batch_end_outputs, batch, batch_idx, dataloader_idx):
+        batch_end_outputs = [opt_idx_out for opt_idx_out in batch_end_outputs if len(opt_idx_out)]
+
+        processed_batch_end_outputs = TrainLoop._prepare_outputs(batch_end_outputs)
+
         # hook
-        self.trainer.call_hook('on_train_batch_end', batch_end_outputs, batch, batch_idx, dataloader_idx)
+        self.trainer.call_hook('on_train_batch_end', processed_batch_end_outputs, batch, batch_idx, dataloader_idx)
         self.trainer.call_hook('on_batch_end')
 
         # figure out what to track for epoch end
@@ -342,6 +346,43 @@ def _process_training_step_output(self, training_step_output, split_batch):
 
         return training_step_output_for_epoch_end, result
 
+    @staticmethod
+    def _prepare_outputs(outputs):
+        """
+        Extract required information from batch or epoch end results.
+        """
+        processed_outputs = []
+        for opt_outputs in outputs:
+            if len(opt_outputs) == 0:
+                continue
+
+            processed_batch_outputs = []
+
+            batch_mode = False
+            if not isinstance(opt_outputs[0], list):
+                opt_outputs = [opt_outputs]
+                batch_mode = True  # these are outputs from batch end
+
+            for batch_outputs in opt_outputs:
+                processed_tbptt_outputs = []
+
+                for tbptt_output in batch_outputs:
+                    out = tbptt_output.extra
+                    out['loss'] = tbptt_output.minimize
+                    processed_tbptt_outputs.append(out)
+
+                if len(processed_tbptt_outputs) == 1:
+                    processed_tbptt_outputs = processed_tbptt_outputs[0]
+                processed_batch_outputs.append(processed_tbptt_outputs)
+
+            if batch_mode:
+                processed_batch_outputs = processed_batch_outputs[0]
+            processed_outputs.append(processed_batch_outputs)
+
+        if len(processed_outputs) == 1:
+            processed_outputs = processed_outputs[0]
+        return processed_outputs
+
     def optimizer_step(self, optimizer, opt_idx, batch_idx, train_step_and_backward_closure):
         model_ref = self.trainer.lightning_module
 
@@ -425,10 +466,15 @@ def run_training_epoch(self):
             if batch_output.signal == -1:
                 break
 
-            batch_end_outputs = self.process_train_step_outputs(batch_output.training_step_output_for_epoch_end)
             # hook
             # TODO: add outputs to batches
-            self.on_train_batch_end(epoch_output, batch_end_outputs, batch, batch_idx, dataloader_idx)
+            self.on_train_batch_end(
+                epoch_output,
+                batch_output.training_step_output_for_epoch_end,
+                batch,
+                batch_idx,
+                dataloader_idx,
+            )
 
             # -----------------------------------------
             # SAVE METRICS TO LOGGERS
@@ -477,11 +523,11 @@ def run_training_epoch(self):
             # progress global step according to grads progress
             self.increment_accumulated_grad_global_step()
 
-        # epoch end hook
+        # handle epoch_output on epoch end
         self.on_train_epoch_end(epoch_output)
 
         # log epoch metrics
-        self.trainer.logger_connector.log_train_epoch_end_metrics(epoch_output, self.num_optimizers)
+        self.trainer.logger_connector.log_train_epoch_end_metrics(epoch_output)
 
         should_check_val = self.should_check_val_fx(batch_idx, is_last_batch, on_epoch=True)
         should_skip_eval = self.trainer.evaluation_loop.should_skip_evaluation(self.trainer.num_val_batches)
@@ -504,6 +550,35 @@ def run_training_epoch(self):
         # progress global step according to grads progress
         self.increment_accumulated_grad_global_step()
 
+    def on_train_epoch_end(self, epoch_output):
+        # prepare epoch output
+        processed_epoch_output = TrainLoop._prepare_outputs(epoch_output)
+
+        model = self.trainer.lightning_module
+
+        if is_overridden('training_epoch_end', model=model):
+            # run training_epoch_end
+            # refresh the result for custom logging at the epoch level
+            model._current_fx_name = 'training_epoch_end'
+
+            # lightningmodule hook
+            training_epoch_end_output = model.training_epoch_end(processed_epoch_output)
+
+            if training_epoch_end_output is not None:
+                raise MisconfigurationException(
+                    'training_epoch_end expects a return of None. '
+                    'HINT: remove the return statement in training_epoch_end'
+                )
+
+            # capture logging
+            self.trainer.logger_connector.cache_logged_metrics()
+
+        # inform logger the batch loop has finished
+        self.trainer.logger_connector.on_train_epoch_end()
+
+        self.trainer.call_hook('on_train_epoch_end', processed_epoch_output)
+        self.trainer.call_hook('on_epoch_end')
+
     def run_training_batch(self, batch, batch_idx, dataloader_idx):
         # track grad norms
         grad_norm_dic = {}
@@ -511,8 +586,10 @@ def run_training_batch(self, batch, batch_idx, dataloader_idx):
         # bookkeeping
         self.trainer.hiddens = None
 
+        optimizers = self.prepare_optimizers()
+
         # track all outputs across time and num of optimizers
-        batch_outputs = [[] for _ in range(len(self.get_optimizers_iterable()))]
+        batch_outputs = [[] for _ in range(len(optimizers))]
 
         if batch is None:
             return AttributeDict(signal=0, grad_norm_dic=grad_norm_dic)
@@ -533,7 +610,7 @@ def run_training_batch(self, batch, batch_idx, dataloader_idx):
         for split_idx, split_batch in enumerate(splits):
 
             # create an iterable for optimizers and loop over them
-            for opt_idx, optimizer in self.prepare_optimizers():
+            for opt_idx, optimizer in optimizers:
 
                 # toggle model params + set info to logger_connector
                 self.run_train_split_start(split_idx, split_batch, opt_idx, optimizer)
@@ -721,13 +798,6 @@ def update_train_loop_lr_schedulers(self, monitor_metrics=None):
             # update lr
             self.trainer.optimizer_connector.update_learning_rates(interval="step", monitor_metrics=monitor_metrics)
 
-    def on_train_epoch_end(self, epoch_output):
-        # inform logger the batch loop has finished
-        self.trainer.logger_connector.on_train_epoch_end()
-
-        self.trainer.call_hook('on_train_epoch_end', epoch_output)
-        self.trainer.call_hook('on_epoch_end')
-
     def increment_accumulated_grad_global_step(self):
         num_accumulated_batches_reached = self._accumulated_batches_reached()
         num_training_batches_reached = self._num_training_batches_reached()
@@ -795,16 +865,6 @@ def save_loggers_on_train_batch_end(self):
         if should_flush_logs and self.trainer.is_global_zero and self.trainer.logger is not None:
             self.trainer.logger.save()
 
-    def process_train_step_outputs(self, all_train_step_outputs):
-        """
-        Figure out what needs to be tracked/logged at the end of the epoch
-        """
-        # the training step outputs a list per optimizer. The list contains the outputs at each time step
-        # when no TBPTT is used, then the list has 1 item per batch
-        # when TBPTT IS used, then the list has n items (1 per time step)
-        # extract one representative sample from each time step (1 if no tbptt) and 0th optimizer
-        return [opt_idx_out for opt_idx_out in all_train_step_outputs if len(opt_idx_out)]
-
     def prepare_optimizers(self):
         # in manual optimization we loop over all optimizers at once
         optimizers = self.get_optimizers_iterable()
diff --git a/tests/callbacks/test_callback_hook_outputs.py b/tests/callbacks/test_callback_hook_outputs.py
index df0eab31aac37..d1bcee43b1f02 100644
--- a/tests/callbacks/test_callback_hook_outputs.py
+++ b/tests/callbacks/test_callback_hook_outputs.py
@@ -26,8 +26,7 @@ def test_train_step_no_return(tmpdir, single_cb: bool):
     class CB(Callback):
 
         def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
-            d = outputs[0][0]
-            assert 'minimize' in d
+            assert 'loss' in outputs
 
         def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
             assert 'x' in outputs
@@ -42,8 +41,7 @@ def on_train_epoch_end(self, trainer, pl_module, outputs):
     class TestModel(BoringModel):
 
         def on_train_batch_end(self, outputs, batch, batch_idx: int, dataloader_idx: int) -> None:
-            d = outputs[0][0]
-            assert 'minimize' in d
+            assert 'loss' in outputs
 
         def on_validation_batch_end(self, outputs, batch, batch_idx: int, dataloader_idx: int) -> None:
             assert 'x' in outputs
@@ -52,8 +50,7 @@ def on_test_batch_end(self, outputs, batch, batch_idx: int, dataloader_idx: int)
             assert 'x' in outputs
 
         def on_train_epoch_end(self, outputs) -> None:
-            d = outputs[0]
-            assert len(d) == self.trainer.num_training_batches
+            assert len(outputs) == self.trainer.num_training_batches
 
     model = TestModel()
 
@@ -79,9 +76,9 @@ class CB(Callback):
 
         def on_validation_epoch_end(self, trainer, pl_module, outputs):
             if trainer.running_sanity_check:
-                assert len(outputs[0]) == trainer.num_sanity_val_batches[0]
+                assert len(outputs) == trainer.num_sanity_val_batches[0]
             else:
-                assert len(outputs[0]) == trainer.num_val_batches[0]
+                assert len(outputs) == trainer.num_val_batches[0]
 
     model = BoringModel()
 
@@ -102,7 +99,7 @@ def test_on_test_epoch_end_outputs(tmpdir):
     class CB(Callback):
 
         def on_test_epoch_end(self, trainer, pl_module, outputs):
-            assert len(outputs[0]) == trainer.num_test_batches[0]
+            assert len(outputs) == trainer.num_test_batches[0]
 
     model = BoringModel()
 
diff --git a/tests/core/test_hooks.py b/tests/core/test_hooks.py
index 191da0a1400c7..087f884d96feb 100644
--- a/tests/core/test_hooks.py
+++ b/tests/core/test_hooks.py
@@ -21,9 +21,9 @@ class TestModel(BoringModel):
 
         def on_validation_epoch_end(self, outputs):
             if trainer.running_sanity_check:
-                assert len(outputs[0]) == trainer.num_sanity_val_batches[0]
+                assert len(outputs) == trainer.num_sanity_val_batches[0]
             else:
-                assert len(outputs[0]) == trainer.num_val_batches[0]
+                assert len(outputs) == trainer.num_val_batches[0]
 
     model = TestModel()
 
@@ -43,7 +43,7 @@ def test_on_test_epoch_end_outputs(tmpdir):
     class TestModel(BoringModel):
 
         def on_test_epoch_end(self, outputs):
-            assert len(outputs[0]) == trainer.num_test_batches[0]
+            assert len(outputs) == trainer.num_test_batches[0]
 
     model = TestModel()
 
diff --git a/tests/models/test_hooks.py b/tests/models/test_hooks.py
index 57af82ccc3e08..93ab34dfa006c 100644
--- a/tests/models/test_hooks.py
+++ b/tests/models/test_hooks.py
@@ -99,7 +99,7 @@ def on_train_epoch_start(self, trainer, pl_module):
             self.len_outputs = 0
 
         def on_train_epoch_end(self, trainer, pl_module, outputs):
-            self.len_outputs = len(outputs[0])
+            self.len_outputs = len(outputs)
 
     class OverriddenModel(BoringModel):
 


Code review comments:
diff_hunk: @@ -342,6 +346,43 @@ def _process_training_step_output(self, training_step_output, split_batch):
 
         return training_step_output_for_epoch_end, result
 
+    @staticmethod
+    def _prepare_outputs(outputs):
path: pytorch_lightning/trainer/training_loop.py
text: @user1:
nit: add typehint - this returns List[Any] or List[List[Any]] ? 

should this take `on_batch` or `on_epoch` as an argument rather than deducing whether we're in batch_mode inside of the function?

diff_hunk: @@ -205,46 +205,40 @@ def log_epoch_metrics_on_evaluation_end(self):
         eval_loop_results = self.trainer.logger_connector.get_evaluate_epoch_results()
         return eval_loop_results
 
-    def __run_eval_epoch_end(self, num_dataloaders):
+    def __run_eval_epoch_end(self, outputs):
         model = self.trainer.lightning_module
 
-        # with a single dataloader don't pass an array
-        outputs = self.outputs
-
-        eval_results = outputs
-        if num_dataloaders == 1:
-            eval_results = outputs[0]
-
+        epoch_end_output = None
         user_reduced = False
 
         if self.trainer.testing:
             if is_overridden('test_epoch_end', model=model):
                 model._current_fx_name = 'test_epoch_end'
-                eval_results = model.test_epoch_end(eval_results)
+                epoch_end_output = model.test_epoch_end(outputs)
                 user_reduced = True
 
         else:
             if is_overridden('validation_epoch_end', model=model):
                 model._current_fx_name = 'validation_epoch_end'
-                eval_results = model.validation_epoch_end(eval_results)
+                epoch_end_output = model.validation_epoch_end(outputs)
                 user_reduced = True
 
         # capture logging
         self.trainer.logger_connector.cache_logged_metrics()
         # depre warning
-        if eval_results is not None and user_reduced:
+        if epoch_end_output is not None and user_reduced:
             step = 'testing_epoch_end' if self.trainer.testing else 'validation_epoch_end'
             self.warning_cache.warn(
                 f'The {step} should not return anything as of 9.1.'
                 ' To log, use self.log(...) or self.write(...) directly in the LightningModule'
             )
 
path: pytorch_lightning/trainer/evaluation_loop.py
text: @user1:
do you know if it's safe to deprecate this? what is 9.1? 

@user2:
Probably a typo for `0.9.1`

Deprecate what exactly?

@user1:
storing the value of `model.validation_epoch_end` or `model.test_epoch_end` and checking if it's not None. like can we delete the code from L235-240? we can do this in a separate PR, just wondering since it seems safe to do so now

@user2:
Ah yes. This is so old it should be fine to just remove it considering we had the warning.

Feel free to do it in a separate PR

@user1:
https://github.com/PyTorchLightning/pytorch-lightning/pull/6973