diff --git a/docs/source/JAX_for_LLM_pretraining.ipynb b/docs/source/JAX_for_LLM_pretraining.ipynb
index 5161aa3..7d4c64c 100644
--- a/docs/source/JAX_for_LLM_pretraining.ipynb
+++ b/docs/source/JAX_for_LLM_pretraining.ipynb
@@ -1013,7 +1013,7 @@
    "id": "3813cbf2",
    "metadata": {},
    "source": [
-    "# Profiling for Hyperparameter Tuning"
+    "## Profiling for Hyperparameter Tuning"
    ]
   },
   {
diff --git a/docs/source/JAX_for_LLM_pretraining.md b/docs/source/JAX_for_LLM_pretraining.md
index a422283..f03d783 100644
--- a/docs/source/JAX_for_LLM_pretraining.md
+++ b/docs/source/JAX_for_LLM_pretraining.md
@@ -462,7 +462,7 @@ checkpointer.save('/content/save', state)
 !ls /content/save/
 ```
 
-# Profiling for Hyperparameter Tuning
+## Profiling for Hyperparameter Tuning
 
 ```{code-cell}
 !pip install -Uq tensorboard-plugin-profile tensorflow tensorboard