diff --git a/nemo/data-flywheel/tool-calling/2_finetuning_and_inference.ipynb b/nemo/data-flywheel/tool-calling/2_finetuning_and_inference.ipynb
index 95a97c75..f3b838cf 100644
--- a/nemo/data-flywheel/tool-calling/2_finetuning_and_inference.ipynb
+++ b/nemo/data-flywheel/tool-calling/2_finetuning_and_inference.ipynb
@@ -562,7 +562,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 56,
+   "execution_count": null,
    "id": "a2e9678c-2785-4e95-b11b-1f41067bc920",
    "metadata": {},
    "outputs": [
@@ -593,8 +593,8 @@
     "    client_with_wandb = nemo_client\n",
     "\n",
     "customization = client_with_wandb.customization.jobs.create(\n",
-    "    name=\"llama-3.2-1b-xlam-ft\",\n",
-    "    output_model=f\"{NMS_NAMESPACE}/llama-3.2-1b-xlam-run1\",\n",
+    "    # Set the finetuned model name in config, So that we can use the same finetuned model name in the all steps in the notebook.\n",
+    "    output_model=f\"{CUSTOM_MODEL}\",\n",
     "    config=f\"{BASE_MODEL}@{BASE_MODEL_VERSION}\",\n",
     "    dataset={\"name\": DATASET_NAME, \"namespace\": NMS_NAMESPACE},\n",
     "    hyperparameters={\n",
@@ -952,6 +952,43 @@
     "print(\"Job Status:\", json.dumps(job_status.model_dump(), indent=2, default=str))"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a64fdf79",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from time import sleep, time\n",
+    "\n",
+    "def wait_job(nemo_client, job_id: str, polling_interval: int = 10, timeout: int = 6000):\n",
+    "    \"\"\"Helper for waiting an eval job using SDK.\"\"\"\n",
+    "    start_time = time()\n",
+    "    job = nemo_client.customization.jobs.retrieve(job_id=job_id)\n",
+    "    status = job.status\n",
+    "\n",
+    "    while (status in [\"pending\", \"created\", \"running\"]):\n",
+    "        # Check for timeout\n",
+    "        if time() - start_time > timeout:\n",
+    "            raise RuntimeError(f\"Took more than {timeout} seconds.\")\n",
+    "\n",
+    "        # Sleep before polling again\n",
+    "        sleep(polling_interval)\n",
+    "\n",
+    "        # Fetch updated status and progress\n",
+    "        job = nemo_client.customization.jobs.retrieve(job_id=job_id)\n",
+    "        status = job.status\n",
+    "\n",
+    "        print(f\"Job status: {status} after {time() - start_time:.2f} seconds.\")\n",
+    "\n",
+    "    return job\n",
+    "\n",
+    "job = wait_job(nemo_client, JOB_ID, polling_interval=5, timeout=2400)\n",
+    "\n",
+    "# Wait for 2 minutes, because sometime, the job is finished, but finetuned model is not ready yet.\n",
+    "sleep(120)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "42b721be-8ca0-4e8f-99a7-5eb12ea1b47f",
diff --git a/nemo/data-flywheel/tool-calling/3_model_evaluation.ipynb b/nemo/data-flywheel/tool-calling/3_model_evaluation.ipynb
index e6c67b30..b14ef038 100644
--- a/nemo/data-flywheel/tool-calling/3_model_evaluation.ipynb
+++ b/nemo/data-flywheel/tool-calling/3_model_evaluation.ipynb
@@ -119,12 +119,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "id": "07faee8b-cc39-4485-abf4-8ab884883d6b",
    "metadata": {},
    "outputs": [],
    "source": [
-    "CUSTOMIZED_MODEL = \"\" # paste from the previous notebook"
+    "CUSTOMIZED_MODEL = CUSTOM_MODEL # paste from the previous notebook"
    ]
   },
   {
diff --git a/nemo/data-flywheel/tool-calling/4_adding_safety_guardrails.ipynb b/nemo/data-flywheel/tool-calling/4_adding_safety_guardrails.ipynb
index 07dd8008..a2a204a7 100644
--- a/nemo/data-flywheel/tool-calling/4_adding_safety_guardrails.ipynb
+++ b/nemo/data-flywheel/tool-calling/4_adding_safety_guardrails.ipynb
@@ -96,12 +96,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "id": "a725db00-7956-466c-8c38-7289fa9fb25e",
    "metadata": {},
    "outputs": [],
    "source": [
-    "CUSTOMIZED_MODEL = \"\" # paste from the previous notebook"
+    "CUSTOMIZED_MODEL = CUSTOM_MODEL # paste from the previous notebook"
    ]
   },
   {
@@ -154,7 +154,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": null,
    "id": "ce022951-1477-48bd-9f53-fc4eaeb43816",
    "metadata": {},
    "outputs": [
@@ -262,7 +262,7 @@
     "# Create guardrails configuration\n",
     "guardrail_config = nemo_client.guardrail.configs.create(\n",
     "    name=\"toolcalling\",\n",
-    "    namespace=\"default\",\n",
+    "    namespace=NMS_NAMESPACE,\n",
     "    data={\n",
     "      \"models\": [\n",
     "        { \n",
@@ -431,7 +431,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
+   "execution_count": null,
    "id": "4fa42b03-3327-4fa9-ab9e-e5bf9bb8b85a",
    "metadata": {},
    "outputs": [],
@@ -454,7 +454,7 @@
     "                }\n",
     "            ],\n",
     "            guardrails={\n",
-    "                \"config_id\": \"toolcalling\"\n",
+    "                \"config_id\": f\"{NMS_NAMESPACE}/toolcalling\"\n",
     "            },\n",
     "            temperature=0.2,\n",
     "            top_p=1\n",
diff --git a/nemo/data-flywheel/tool-calling/config.py b/nemo/data-flywheel/tool-calling/config.py
index 337625ed..4ae2bad4 100644
--- a/nemo/data-flywheel/tool-calling/config.py
+++ b/nemo/data-flywheel/tool-calling/config.py
@@ -19,3 +19,6 @@
 # (Optional) Configure the base model. Must be one supported by the NeMo Customizer deployment!
 BASE_MODEL = "meta/llama-3.2-1b-instruct"
 BASE_MODEL_VERSION = "v1.0.0+A100"
+
+# (Optional) Configure the finetuned model name. So that we can use the same finetuned model name in the all steps in the notebook.
+CUSTOM_MODEL = f"{NMS_NAMESPACE}/llama-3.2-1b-xlam-run1@v1"