import torch


def _convert_dtype(dtype):
    """Transformer4Rec uses these fixed dtypes and this function converts the original dtype
    to this fixed dtypes"""
    if dtype in [torch.float64, torch.float32, torch.float16]:
        return torch.float32
    if dtype in [
        torch.int64,
        torch.int32,
        torch.int16,
        torch.int8,
        torch.uint8,
    ]:
        return torch.long

    raise ValueError(f"Can't convert dtype {dtype})")