def evaluate_checkpoint(params, batch):
    return [[float(compute_loss(params, sample[None, :]))] for sample in batch]

def aggregate_losses(outputs):
    losses = [row[0] for batch in outputs for row in batch]
    return sum(losses) / len(losses)