import torch


def get_elem_size(tensor_obj):
    """

    :param tensor_obj: a tensor
    :return: the size in bytes of a single element in the tensor based on its
    type

    >>> t1 = tensor([1,2,3], dtype=torch.int32)
    >>> size = get_elem_size(t1)
    >>> assert size == 4
    >>> del t1
    >>> t2 = tensor([1.,2.,3.], dtype=torch.float)
    >>> size = get_elem_size(t2)
    >>> assert size == 4
    >>> del t2
    >>> t3 = tensor([1.,2.,3.], dtype=torch.float16)
    >>> size = get_elem_size(t3)
    >>> assert size == 2
    >>> del t3
    >>> t4 = tensor([1.,2.,3.], dtype=torch.double)
    >>> size = get_elem_size(t4)
    >>> assert size == 8
    >>> del t4
    """
    obj_type = tensor_obj.dtype

    if (obj_type is torch.float64) or (obj_type is torch.double) or (
            obj_type is torch.int64) or (obj_type is torch.long):
        return 8
    elif (obj_type is torch.float32) or (obj_type is torch.float) or (
            obj_type is torch.int32) or (obj_type is torch.int):
        return 4
    elif (obj_type is torch.float16) or (obj_type is torch.half) or (
            obj_type is torch.int16) or (obj_type is torch.short):
        return 2
    elif (obj_type is torch.int8) or (obj_type is torch.uint8):
        return 1
    else:
        raise AttributeError(f"Unknown torch type: {obj_type}")