@measurement.evaluate(params=0, batch=1)
def evaluate(params, batch):
    features = batch[:, :10].to(dtype=params["weight"].dtype)
    return functional_call(model, params, (features,))

@measurement.aggregate
def aggregate(outputs):
    values = [
        value
        for output_batch in outputs
        for row in output_batch
        for value in row
    ]
    return sum(values) / len(values)