diff --git a/docs/notebooks/train_transformer.ipynb b/docs/notebooks/train_transformer.ipynb
index dc3a496..c3bcb08 100644
--- a/docs/notebooks/train_transformer.ipynb
+++ b/docs/notebooks/train_transformer.ipynb
@@ -427,7 +427,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Load the dataset"
+    "### Load the dataset"
    ]
   },
   {
@@ -448,7 +448,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Tokenizer"
+    "### Tokenizer"
    ]
   },
   {
@@ -489,7 +489,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Dataloader"
+    "### Dataloader"
    ]
   },
   {
@@ -946,7 +946,11 @@
     "    text_en = test_dataset[index][\"translation\"][\"en\"]\n",
     "    text_en_pred = translate_from_arabic_to_english(text_ar, key=jr.PRNGKey(0))\n",
     "\n",
-    "    print(f\"input arabic: {text_ar}\\n\" f\"true english: {text_en}\\n\" f\"pred english: {text_en_pred}\\n\")"
+    "    print(\n",
+    "        f\"input arabic: {text_ar}\\n\"\n",
+    "        f\"true english: {text_en}\\n\"\n",
+    "        f\"pred english: {text_en_pred}\\n\"\n",
+    "    )"
    ]
   }
  ],