transformers

49303a2d - Address @bastings comment in PR, moving jax.jit decorated outside of __call__