def _jit_predict_fn(model_predict, metric_fn, n_devices, jit=True):
  """Returns a JIT-compiled predict function (unless jit=False)."""
  model = tl.Serial(model_predict, metric_fn)
  if not jit:
    return model.pure_fn

  return tl.jit_forward(model.pure_fn, n_devices)
Returns a JIT-compiled predict function (unless jit=False).
trax/supervised/trainer_lib.py