jax.jit training_util.py::make_predictions.

PiperOrigin-RevId: 617565334
tensorflow · Mar 20, 2024 · f7f04c8 · f7f04c8
1 parent b0abbc7
commit f7f04c8
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tensorflow_probability/python/experimental/autobnn/training_util.py b/tensorflow_probability/python/experimental/autobnn/training_util.py
@@ -211,7 +211,8 @@ def _plot_loss_fn(losses, ax=None, log_scale=True) -> plt.Figure:
 
 def make_predictions(params, net: bnn.BNN, x_test: jax.Array) -> jax.Array:
   """Use a (batch of) parameters to make a prediction on x_test data."""
-  return jax.vmap(lambda p: net.apply(p, x_test))(params)
+  f = jax.jit(lambda p: net.apply(p, x_test))
+  return jax.vmap(f)(params)
 
 
 def make_results_dataframe(