import torch


def _get_device(device_string):
    """Returns a `torch.device`

    Args:
      device_string: 'cpu' or 'cuda'

    Returns:
      `torch.device`
    """
    if device_string.lower() == 'cpu':
      return torch.device('cpu')
    if device_string.lower() == 'cuda':
      if torch.cuda.device_count() == 0:
        print("Warning: There's no GPU available on this machine!")
        return None
      return torch.device('cuda:0')
    raise Exception(
      '{} is not a valid option. Choose `cpu` or `cuda`.'.format(device_string))