def METHOD_NAME(device: torch.device) -> torch.device:
    if device.type == "cuda" and device.index is None:
        return torch.device("cuda", index=torch.cuda.current_device())
    if device.type == "cpu" and device.index is not None:
        return torch.device("cpu")
    return device