training_epoch_end log output gets combined with next epoch training
🐛 Bug
So, I put 'training_epoch_end' function in my LightningModule. I have it return this dictionary
{'log':{'train_loss': tensor(0.3616, device='cuda:0'), 'epoch': 0}
I check the run_training_epoch_end function in the Pytorch library, it looks like it is working normally as log_epoch_metrics is showing the 'log' part in the dictionary produced by 'training_epoch_end' function
{'train_loss': tensor(0.3616, device='cuda:0'), 'epoch': 0}
So, they send it off to the logger. But there is problem, it is trying to combine the dictionary above with the results from the training step fo the next epoch. When I check the variable self._metrics_to_agg, I get the following result. Of course, it is impossible to combine these dictionaries as they have different keys. I guess the main problem is that the code is combining the log results of run_training_epoch_end function with the results of the next training batch.
[{'train_loss': tensor(0.3616, device='cuda:0'), 'epoch': 0}, {'loss': 0.48756, 'input': ..., 'ouput': ...}]
Any ideas to solve this problem? I will appreciate your help! Here is the whole error stack:
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-5-a4569c0ab5cd> in <module>
      3 logging.getLogger('allennlp').setLevel(logging.INFO)
      4 
----> 5 train("/shared/labs/workflow/configs/ner_span_path_2.jsonnet", "/Dropbox/EvidLabs/atee/lab_data_enrichment_project/train/DataEnrichmentSpanPath/DataEnrichmentSpanPath__INV_M_span_path_42_43", AllenNlpLightningModule, tags=[])
      6 
      7 # train("/shared/labs/workflow/configs/candidate_span_grouper.jsonnet", "./DataEnrichmentBunchedCrfTaggerTest", AllenNlpLightningModule, tags=["Frozen"])

/code/evid-research/evid2/pytorch_lightning/trainer/config_file_trainer.py in train(config_file, project_dir, lightning_module, run_id, run_name_override, tags, dry_run)
    171     # 3 START TRAINING
    172     # ------------------------
--> 173     trainer.fit(model)
    174 
    175 

/code/pytorch-lightning/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders)
    971         # easier to avoid NCCL issues
    972         elif self.use_dp:
--> 973             self.dp_train(model)
    974 
    975         elif self.use_horovod:

/code/pytorch-lightning/pytorch_lightning/trainer/distrib_parts.py in dp_train(self, model)
    265         model = LightningDataParallel(model, device_ids=device_ids)
    266 
--> 267         self.run_pretrain_routine(model)
    268 
    269         model.forward = model_autocast_original_forward

/code/pytorch-lightning/pytorch_lightning/trainer/trainer.py in run_pretrain_routine(self, model)
   1154 
   1155         # CORE TRAINING LOOP
-> 1156         self.train()
   1157 
   1158     def test(

/code/pytorch-lightning/pytorch_lightning/trainer/training_loop.py in train(self)
    368                 # RUN TNG EPOCH
    369                 # -----------------
--> 370                 self.run_training_epoch()
    371 
    372                 if self.max_steps and self.max_steps <= self.global_step:

/code/pytorch-lightning/pytorch_lightning/trainer/training_loop.py in run_training_epoch(self)
    480             # SAVE METRICS TO LOGGERS
    481             # -----------------------------------------
--> 482             self.save_train_loop_metrics_to_loggers(batch_idx, batch_output)
    483 
    484             # progress global step according to grads progress

/code/pytorch-lightning/pytorch_lightning/trainer/training_loop.py in save_train_loop_metrics_to_loggers(self, batch_idx, batch_output)
    564             # logs user requested information to logger
    565             print("batch_output.batch_log_metrics", batch_output.batch_log_metrics.keys())
--> 566             self.log_metrics(batch_output.batch_log_metrics, batch_output.grad_norm_dic)
    567 
    568     def save_loggers_in_training_loop(self, batch_idx):

/code/pytorch-lightning/pytorch_lightning/trainer/logging.py in log_metrics(self, metrics, grad_norm_dic, step)
     72         if self.is_global_zero and self.logger is not None:
     73             print("scalar_metrics", scalar_metrics.keys())
---> 74             self.logger.agg_and_log_metrics(scalar_metrics, step=step)
     75             self.logger.save()
     76 

/code/pytorch-lightning/pytorch_lightning/loggers/base.py in agg_and_log_metrics(self, metrics, step)
    131         """
    132         print("metrics", metrics.keys())
--> 133         agg_step, metrics_to_log = self._aggregate_metrics(metrics=metrics, step=step)
    134 
    135         if metrics_to_log:

/code/pytorch-lightning/pytorch_lightning/loggers/base.py in _aggregate_metrics(self, metrics, step)
     88 
     89         # compute the metrics
---> 90         agg_step, agg_mets = self._reduce_agg_metrics()
     91 
     92         # as new step received reset accumulator

/code/pytorch-lightning/pytorch_lightning/loggers/base.py in _reduce_agg_metrics(self)
    109             print("self._agg_key_funcs", self._agg_key_funcs)
    110             print("self._agg_default_func", self._agg_default_func)
--> 111             agg_mets = merge_dicts(self._metrics_to_agg, self._agg_key_funcs, self._agg_default_func)
    112         return self._prev_step, agg_mets
    113 

/code/pytorch-lightning/pytorch_lightning/loggers/base.py in merge_dicts(dicts, agg_key_funcs, default_func)
    386             print("fn or default_func", fn or default_func)
    387             print("values_to_agg", values_to_agg)
--> 388             d_out[k] = (fn or default_func)(values_to_agg)
    389 
    390     return d_out

<__array_function__ internals> in mean(*args, **kwargs)

/venvs/dev3.7/lib/python3.7/site-packages/numpy/core/fromnumeric.py in mean(a, axis, dtype, out, keepdims)
   3333 
   3334     return _methods._mean(a, axis=axis, dtype=dtype,
-> 3335                           out=out, **kwargs)
   3336 
   3337 

/venvs/dev3.7/lib/python3.7/site-packages/numpy/core/_methods.py in _mean(a, axis, dtype, out, keepdims)
    149             is_float16_result = True
    150 
--> 151     ret = umr_sum(arr, axis, dtype, out, keepdims)
    152     if isinstance(ret, mu.ndarray):
    153         ret = um.true_divide(

TypeError: unsupported operand type(s) for +: 'dict' and 'dict'

To Reproduce
Steps to reproduce the behavior:

Add a training_epoch_end function to your Lightning Module and run it. You can use mine in the code sample section if you want. The key is that the "log"  section of training_epoch_end dictionary output must be of a different format than the dictionary that contains the results of your training sample while training (like in the example I provided {'train_loss': tensor(0.3616, device='cuda:0'), 'epoch': 0} and {'loss': 0.48756, 'input': ..., 'ouput': ...} are different formats as they don't share the same keys).

Code sample
 def training_epoch_end(self, outputs):
        print("training_epoch_end")
        print(len(outputs))

        prefix = "train_"
        metric_modules = self.training_metric_modules

        # Handle loss:
        avg_loss = torch.stack([x['loss'] for x in outputs]).mean()
        result = {'log': {}, 'progress_bar': {}}
        result[prefix + 'loss'] = avg_loss
        result['log'][prefix + 'loss'] = avg_loss
        result['progress_bar'][prefix + 'loss'] = avg_loss

        # Add tracking variables
        result['log']['epoch'] = self.current_epoch

        return result

Expected behavior
I thought I would be able to run training_epoch_end function with no combination with the training samples in the next epoch. I expected no error, like running validation_epoch_end.
Environment

CUDA:
- GPU:
- Tesla V100-SXM2-16GB
- available:         True
- version:           10.2
Packages:
- numpy:             1.18.5
- pyTorch_debug:     False
- pyTorch_version:   1.5.0
- pytorch-lightning: 0.8.4
- tensorboard:       2.2.2
- tqdm:              4.46.1
System:
- OS:                Linux
- architecture:
- 64bit
-
- processor:         x86_64
- python:            3.7.7
- version:           #113-Ubuntu SMP Wed Jan 29 14:54:54 UTC 2020