from typing import Union
import torch
from typing import Sequence
from typing import Iterator


def count_tensor_list_size(tensor_list: Union[torch.Tensor, Sequence, Iterator], format="byte"):
    """Approximately count the size of a list of tensors in terms of bytes."""
    if torch.is_tensor(tensor_list):
        tensor_list = [tensor_list]
    _bytes = 4 * sum([t.numel() for t in tensor_list])
    if format == "byte":
        return _bytes
    elif format == "kb":
        return _bytes / 1024
    elif format == "mb":
        return _bytes / 1024 ** 2
    elif format == "gb":
        return _bytes / 1024 ** 3
    else:
        raise ValueError(f"Unknown format: {format}")