import torch


def init_torch_device(gpu_id: int=-1) -> torch.device:
    """ Creates torch device based on the GPU ID

    Args:
        gpu_id (int, optional): CUDA device index, CPU if < 0
    """
    device_str = 'cpu'
    if torch.cuda.is_available():
        num_gpus = torch.cuda.device_count()
        if gpu_id >= 0 and gpu_id < num_gpus:
            device_str = f'cuda:{gpu_id}'

    return torch.device(device_str)