jitted predict

[soldierofhell]

Hi, I'm starting to explore your framework. I'm familiar with jax, but not with objax. I noticed that train ops are jitted with objax.Jit, but as my goal is to have fast prediction embedded in some larger jax code, I wonder if predit() can be also jitted? Thanks in advance,

Regards,

[wil-j-wil]

Yes you can jit compile that part just like any other, just include it in your larger chunk of jitted code.
Or if you want to jit that part on its own you can do

predict = objax.Jit(model.predict, model.vars() + opt_hypers.vars())

[wil-j-wil]

Assuming the above solves this issue. If not, feel free to re-open.