def get_output_selected_model(neuron_pos, session_id, model):
    """Creates a version of the model that has its output selected down to a single uniquely identified neuron.

    Args:
        neuron_pos: An integer, the position of the neuron in the model's output.
        session_id: A string that uniquely identifies one of the model's readouts.
        model: A PyTorch module that can be called with a keyword argument called "data_key". The output of the
            module is expected to be a two dimensional Torch tensor where the first dimension corresponds to the
            batch size and the second to the number of neurons.

    Returns:
        A function that takes the model input(s) as parameter(s) and returns the model output corresponding to the
        selected neuron.
    """

    def output_selected_model(x, *args, **kwargs):
        output = model(x, *args, data_key=session_id, **kwargs)
        return output[:, neuron_pos]

    return output_selected_model