from typing import Optional
import torch


def get_n_byte_tensor(n: int, device: Optional[torch.device] = None) -> torch.Tensor:
    """Returns a torch.int8 tensor of size n.

    Args:
        n: size of the tensor to allocate
        device: torch.device to allocate tensor on

    Returns:
        torch.int8 tensor of size n on device

    Notes: Though 1 int8 = 1 byte, memory is allocated in blocks, such that the size
    of the tensor in bytes >= n.
    """
    return torch.zeros(n, dtype=torch.int8, device=device)