Speed up training by making use of `jax.jit`

[phinate]

The working CLs training loop has so far not made use of the @jax.jit decorator. Doing so would dramatically speed up training time.