diff --git a/notebook/oai_completion.ipynb b/notebook/oai_completion.ipynb
index f25db693fb81..d3aaacb17cb8 100644
--- a/notebook/oai_completion.ipynb
+++ b/notebook/oai_completion.ipynb
@@ -404,6 +404,8 @@
     "Users can specify tuning data, optimization metric, optimization mode, evaluation function, search spaces, etc. The default search space is:\n",
     "\n",
     "```python\n",
+    "from flaml import tune\n",
+    "\n",
     "default_search_space = {\n",
     "    \"model\": tune.choice([\n",
     "        \"text-ada-001\",\n",