def benchmark_torch_function(iters: int, function, *args, **kwargs) -> float:
    """
    function for benchmarking a pytorch function.

    Parameters
    ----------
    iters: int
        Number of iterations.
    function: lambda function
        function to benchmark.
    args: Any type
        Args to function.

    Returns
    -------
    float
        Runtime per iteration in ms.
    """
    import torch

    # Warm up
    for _ in range(5):
        function(*args, **kwargs)

    # Start benchmark.
    torch.cuda.synchronize()
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    start_event.record()
    for _ in range(iters):
        function(*args, **kwargs)
    end_event.record()
    torch.cuda.synchronize()
    # in ms
    return (start_event.elapsed_time(end_event)) / iters

    function for benchmarking a pytorch function.

    Parameters
    ----------
    iters: int
        Number of iterations.
    function: lambda function
        function to benchmark.
    args: Any type
        Args to function.

    Returns
    -------
    float
        Runtime per iteration in ms.
    
python/aitemplate/testing/benchmark_pt.py