diff --git a/agent_app_sample_code/02_data_pipeline.py b/agent_app_sample_code/02_data_pipeline.py
index 53d0918..b8a0570 100644
--- a/agent_app_sample_code/02_data_pipeline.py
+++ b/agent_app_sample_code/02_data_pipeline.py
@@ -97,9 +97,7 @@
 
 # Embedding model endpoint. The list of off-the-shelf embeddings can be found here:
 # https://docs.databricks.com/en/machine-learning/foundation-models/index.html
-EMBEDDING_MODEL_ENDPOINT = "databricks-bge-large-en"
-# EMBEDDING_MODEL_ENDPOINT = "ep-embeddings-small"
-# EMBEDDING_MODEL_ENDPOINT = "bge-test"
+EMBEDDING_MODEL_ENDPOINT = "databricks-gte-large-en"
 
 # COMMAND ----------