your_name/torch_example

@measurement.evaluate(params=0, batch=1)
def evaluate(params, batch):
    features = batch[:, :10].to(dtype=params["weight"].dtype)
    # functional_call evaluates the module with the measured checkpoint while
    # leaving the model object unchanged between batches.
    return functional_call(model, params, (features,))

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