Perhaps separating the function for building the model and running the inference.
Perhaps separating the function for building the model and running the inference.