def run_module_with_pools(
    *,
    ait_module,
    inputs_pool,
    outputs_pool,
    num_iters,
    stream_ptr: Optional[int] = None,
    sync: bool = False,
    graph_mode: bool = False,
):
    """
    Run the module with the given inputs and outputs pools.
    Parameters
    ----------
    ait_module : Model
        The AIT module to run.
    inputs_pool : List[Dict[str, torch.Tensor]]
        A list of inputs to pass into Model.RunWithTensors.
    outputs_pool : List[Dict[str, torch.Tensor]]
        A list of outputs to pass into Model.RunWithTensors.
    num_iters : int
        The number of iterations to run.
    stream_ptr : Optional[int]
        The CUDA stream pointer to run the module on; if None, use the legacy stream.
    sync : bool
        Whether to synchronize the CUDA stream after each iteration.
    graph_mode : bool
        Whether to run the module in graph mode.
    """
    for i in range(num_iters):
        ait_module.run_with_tensors(
            inputs_pool[i % len(inputs_pool)],
            outputs_pool[i % len(outputs_pool)],
            sync=sync,
            stream_ptr=stream_ptr,
            graph_mode=graph_mode,
        )
Run the module with the given inputs and outputs pools.
Parameters
----------
ait_module : Model
The AIT module to run.
inputs_pool : List[Dict[str, torch.Tensor]]
A list of inputs to pass into Model.RunWithTensors.
outputs_pool : List[Dict[str, torch.Tensor]]
A list of outputs to pass into Model.RunWithTensors.
num_iters : int
The number of iterations to run.
stream_ptr : Optional[int]
The CUDA stream pointer to run the module on; if None, use the legacy stream.
sync : bool
Whether to synchronize the CUDA stream after each iteration.
graph_mode : bool
Whether to run the module in graph mode.
python/aitemplate/testing/benchmark_ait.py