diff --git a/course/en/chapter3/section4.ipynb b/course/en/chapter3/section4.ipynb
index 397b34d1..f2ed130f 100644
--- a/course/en/chapter3/section4.ipynb
+++ b/course/en/chapter3/section4.ipynb
@@ -148,7 +148,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from transformers import AdamW\n",
+    "from torch.optim import AdamW\n",
     "\n",
     "optimizer = AdamW(model.parameters(), lr=5e-5)"
    ]
@@ -270,7 +270,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from transformers import AdamW, AutoModelForSequenceClassification, get_scheduler\n",
+    "from transformers import AutoModelForSequenceClassification, get_scheduler\n",
+    "from torch.optim import AdamW\n",
     "\n",
     "model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)\n",
     "optimizer = AdamW(model.parameters(), lr=3e-5)\n",
@@ -310,7 +311,8 @@
    "outputs": [],
    "source": [
     "from accelerate import Accelerator\n",
-    "from transformers import AdamW, AutoModelForSequenceClassification, get_scheduler\n",
+    "from transformers import AutoModelForSequenceClassification, get_scheduler\n",
+    "from torch.optim import AdamW\n",
     "\n",
     "accelerator = Accelerator()\n",
     "\n",