from typing import List
import torch
from typing import Tuple
import time


def compute_torch_latency(
    xs: List[torch.Tensor],
    model: torch.nn.Module,
    device: str,
    steps: int,
) -> Tuple[float, List[float]]:
    """Compute the latency associated with the torch model.

    Args:
        xs (List[Tensor]): List of input tensors (a single batch for the model)
        model (Module): Torch model.
        device (str): Device where computing the latency.
        steps (int): Number of times the experiment needs to be performed for
            computing the statistics.

    Returns:
        Float: Average latency.
        List[Float]: List of latencies obtained.
    """
    xs = [x.to(device) for x in xs]
    model = model.to(device)
    latencies = []
    for _ in range(steps):
        starting_time = time.time()
        _ = model.forward(*xs)
        latencies.append(time.time() - starting_time)
    latency = sum(latencies) / steps
    return latency, latencies