def tensor_size_bytes(tensor, unit='MB'):
    """
    Get the size of the tensor in bytes, or a unit that's multiple of bytes
    :param tensor: the pytorch tensor
    :param unit: GigaBytes or GB (assumes GB=1e9 Bytes), MegaBytes or MB (assumes MB=1e6 Bytes),
    KiloBytes or KB (assumes KB=1e3 Bytes), Bytes or B
    :return: the size of the tensor in the desired unit
    """
    if 'G' in unit.upper():
        factor = 1e9
    elif 'M' in unit.upper():
        factor = 1e6
    elif 'K' in unit.upper():
        factor = 1e3
    else:
        factor = 1.0
    return (tensor.element_size() * tensor.nelement()) / factor