Weight batch losses by the number of examples

An average of batch averages gives each batch equal weight. If you want each example to have equal weight, account for how many examples each batch contains.

A short batch changes the weighting

batch 1: 8 examples, mean loss 1.0
batch 2: 2 examples, mean loss 5.0

average of batch means: (1.0 + 5.0) / 2 = 3.0
average over examples:  (8 * 1.0 + 2 * 5.0) / 10 = 1.8

The unweighted calculation treats the two-example batch as equally important as the eight-example batch. Multiplying each mean by its count recovers its loss total, so the combined mean is:

sum(batch_mean * batch_size) / sum(batch_size)

Equal batch sizes guarantee that the two formulas agree mathematically. They can also agree with unequal sizes, for example when all batch means are identical. Unequal sizes make an unweighted average unreliable, not necessarily different on every input.

Use the count that matches the mean

This exercise assumes each batch mean averages one loss per example. Under that contract, batch size is the correct weight.

Other metrics may average over valid tokens, pixels or a sum of sample weights. Use that denominator instead. A loss with reduction="sum" can supply the numerator directly, but the final denominator must still match the quantity you intend to average.

For reporting metrics in a real training loop, accumulate detached values so logging does not retain computation graphs. Keep accumulator tensors on a compatible device and use an appropriate dtype.

When batches have different sizes

A final batch can be shorter when the number of sampled examples is not divisible by the batch size. It is not shorter in every epoch or every loader configuration. Dropping it changes which examples are included; it does not compute a metric over the omitted examples.

The same weighting principle applies to batch accuracies and metrics combined across workers. Sum the matching numerators and denominators rather than averaging local means when their counts differ. Unequal micro-batches also need appropriate weighting during gradient accumulation, though the training framework's gradient averaging must be accounted for.

def epoch_loss(batch_means: list, batch_sizes: list) -> torch.Tensor