from typing import Union
import torch
import logging


def get_device(
    device: Union[None, str, torch.device],
) -> torch.device:
    """Resolve the device, either specified as name, or device."""
    if device is None:
        device = 'cuda'
    if isinstance(device, str):
        device = torch.device(device=device)
    assert isinstance(device, torch.device)
    if not torch.cuda.is_available() and device.type == 'cuda':
        logging.warning(f'Requested device {device}, but CUDA is unavailable. Falling back to cpu.')
        device = torch.device('cpu')
    return device