def load_classifier_func(
    key: jnp.ndarray,
    sample: Dict,
    image_keys: List[str],
    checkpoint_path: str,
    step: Optional[int] = None,
) -> Callable[[Dict], jnp.ndarray]:
    """
    Return: a function that takes in an observation
            and returns the logits of the classifier.
    """
    classifier = create_classifier(key, sample, image_keys)
    classifier = checkpoints.restore_checkpoint(
        checkpoint_path,
        target=classifier,
        step=step,
    )
    func = lambda obs: classifier.apply_fn(
        {"params": classifier.params}, obs, train=False
    )
    func = jax.jit(func)
    return func
Return: a function that takes in an observation
and returns the logits of the classifier.
serl_launcher/serl_launcher/networks/reward_classifier.py