diff --git a/course/en/chapter3/section2.ipynb b/course/en/chapter3/section2.ipynb
index 2359cbe78..f03ee3a5f 100644
--- a/course/en/chapter3/section2.ipynb
+++ b/course/en/chapter3/section2.ipynb
@@ -30,7 +30,8 @@
    "outputs": [],
    "source": [
     "import torch\n",
-    "from transformers import AdamW, AutoTokenizer, AutoModelForSequenceClassification\n",
+    "from torch.optim import AdamW\n",
+    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
     "\n",
     "# Same as before\n",
     "checkpoint = \"bert-base-uncased\"\n",