def _video_embedding(checkpoint_path: str):
    'Load the video embedding for the BraVe model to evaluate.'
    checkpoint = np.load(checkpoint_path, allow_pickle=True).item()
    params = checkpoint['params']
    state = checkpoint['state']
    brave_config_dct = checkpoint['config']
    brave_config = brave.BraveConfig(**brave_config_dct)
    model = brave.get_model(brave_config)

    @jax.jit
    def embedding_fn(view: datasets.View) -> chex.Array:
        narrow_forward_fn = model.forward_fns['narrow_video']
        (embedding, _) = narrow_forward_fn(params, state, None, view, False)
        return embedding

    def synchronous_embedding_fn(view: datasets.View) -> chex.Array:
        return jax.device_get(embedding_fn(view))
    return synchronous_embedding_fn
Load the video embedding for the BraVe model to evaluate.
def _video_embedding(checkpoint_path: str):
    
    checkpoint = np.load(checkpoint_path, allow_pickle=True).item()
    params = checkpoint['params']
    state = checkpoint['state']
    brave_config_dct = checkpoint['config']
    brave_config = brave.BraveConfig(**brave_config_dct)
    model = brave.get_model(brave_config)

    @jax.jit
    def embedding_fn(view: datasets.View) -> chex.Array:
        narrow_forward_fn = model.forward_fns['narrow_video']
        (embedding, _) = narrow_forward_fn(params, state, None, view, False)
        return embedding

    def synchronous_embedding_fn(view: datasets.View) -> chex.Array:
        return jax.device_get(embedding_fn(view))
    return synchronous_embedding_fn