def device_mapping(cuda_device: int) ->str:
    """
    In order to `torch.load()` a GPU-trained model onto a CPU (or specific GPU),
    you have to supply a `map_location` function. Call this with
    the desired `cuda_device` to get the function that `torch.load()` needs.
    """
    if cuda_device < 0:
        return 'cpu'
    return f'cuda:{cuda_device}'
