Per-epoch loss computation over just the epoch instead of snapshot of total average

[largestack]

Quick request if possible. loss/epoch seems to be calculated as the total average loss since the start of training. Is it possible to upgrade this to record the average step loss of the epoch instead? I think this will better-guide everyone in how their latest model is performing and better matches what people expect this value to contain.

Eg, something like:

for epoch in range(epoch_to_start, num_train_epochs):
    epoch_losses = []  # Reset each epoch
    for step, batch in enumerate(train_dataloader):
        # ... training loop ...
        current_loss = loss.detach().item()
        epoch_losses.append(current_loss)
        # ... other code ...

    # After steps in epoch, calculate epoch loss
    avg_epoch_loss = sum(epoch_losses) / len(epoch_losses) if len(epoch_losses) > 0 else 0.0
    if len(accelerator.trackers) > 0:
        logs = {"loss/epoch": avg_epoch_loss}
        accelerator.log(logs, step=epoch + 1)