import torch


def make_device(device_str: str):
    """
    return device from string that can be either 'cpu', 'cuda' or cuda device number"
    """
    try:
        return torch.device(int(device_str))  # Can fail on conversion to int
    except ValueError:
        return torch.device(device_str)