diff --git a/docs/modules/prompts/prompt_templates/examples/connecting_to_a_feature_store.ipynb b/docs/modules/prompts/prompt_templates/examples/connecting_to_a_feature_store.ipynb
index da7d313d..ae58fbfd 100644
--- a/docs/modules/prompts/prompt_templates/examples/connecting_to_a_feature_store.ipynb
+++ b/docs/modules/prompts/prompt_templates/examples/connecting_to_a_feature_store.ipynb
@@ -110,7 +110,7 @@
     "                'driver_hourly_stats:acc_rate',\n",
     "                'driver_hourly_stats:avg_daily_trips'\n",
     "            ],\n",
-    "            entity_rows=[{\"driver_id\": 1001}]\n",
+    "            entity_rows=[{\"driver_id\": driver_id}]\n",
     "        ).to_dict()\n",
     "        kwargs[\"conv_rate\"] = feature_vector[\"conv_rate\"][0]\n",
     "        kwargs[\"acc_rate\"] = feature_vector[\"acc_rate\"][0]\n",