@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
@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