The mean of the means

The standard training loop reports an epoch loss:

losses = []
for batch in loader:
    loss = criterion(model(batch.x), batch.y)      # already a mean
    losses.append(loss)
epoch_loss = torch.stack(losses).mean()

This is the mean of per-batch means, and it is only the mean over the dataset when every batch is the same size.

The arithmetic

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

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

The two-example batch gets the same vote as the eight-example one. With drop_last=False, which is the default, every epoch ends with a short batch, so this is not an edge case: it happens once per epoch, every epoch.

How wrong it is depends on how much the short batch differs, which is exactly the situation where you were relying on the number.

The fix

Weight by the count:

total = torch.zeros(())
count = 0
for batch in loader:
    n = batch.x.size(0)
    total += criterion(...) * n
    count += n
epoch_loss = total / count

Or keep reduction="sum" on the criterion and divide once at the end, which says the same thing with less bookkeeping.

Where else it bites

• Accuracy computed as the mean of per-batch accuracies. Same bug.
• Gradient accumulation, where the micro-batches have different sizes: averaging the losses weights them wrongly, so the gradient is not the gradient of the full batch.
• Distributed training, where each rank may hold a different number of examples and averaging the per-rank losses is subtly off.

drop_last=True sidesteps all of it by discarding the short batch, which is common for training and wrong for evaluation, where you would be silently scoring on a subset.

Your task

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

Given each batch’s mean loss and its size, return the mean over all examples.