Batch weight for accumulating gradients
🚀 Feature
Support different weights for each batch during batch gradient accumulation.
Motivation
In many cases not every batch are equal. For example, when the inputs are sequences of tokens of variable length, the mean loss of several batches is not the average across the individual batch losses, since each batch would have varying number of samples.
Pitch
In training_loop.py, add a 'batch_total_weight' attribute initialized as zero. When fetching returns from training_step, if a certain 'batch_weight' key is in the returning dictionary, the batch's weight would be multiplied by 'batch_weight' during backward step and 'batch_weight' will be added to 'batch_total_weight'.  Before optimization step, the gradients are scaled by 'batch_total_weight' and the latter would be set to zero for next accumulating steps.
A typical scenario is when 'batch_weight' is the number of samples within each batch, so that the average across multiple batches are computed correctly.