Compute Loss on Accumulated Batches
🚀 Feature
Make it possible to accumulate the outputs and targets from the training step for the number of grad batches that should be accumulated. Then compute the loss over the accumulated outputs and targets.
Motivation
The accumulate grad batches option is great for people who only have access to a GPU with limited memory. However, some loss functions are not batch independent, i.e. the loss of one sample depends on the other samples within a batch. For example the F1 score or Matthew's correlation coefficient can be modified so they can be directly applied as loss functions. In these cases, it would be great to be able to accumulate the outputs and compute the loss afterwards, similarly to how it works in validation_end and test_end, except before the backwards step.
Pitch
Add an accumulate function that takes the accumulated output of the training_step. The default behavior would be to output the sum over all keys in the dict and to normalize by the number of accumulated grad batches. This can then be overwritten though to compute the loss over the entire accumulated batch. This would not break the current API but accommodate batch dependent loss functions.