from typing import Optional
import torch


def get_device(device: Optional[str] = None) -> torch.device:
    """Selects a device to perform an NST on.

    Args:
        device: If ``str``, returns the corresponding :class:`~torch.device`. If
            ``None`` selects CUDA if available and otherwise CPU. Defaults to ``None``.
    """
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
    return torch.device(device)