diff --git a/ch03/02_bonus_efficient-multihead-attention/mha-implementations.ipynb b/ch03/02_bonus_efficient-multihead-attention/mha-implementations.ipynb
index fb490da..a1d074b 100644
--- a/ch03/02_bonus_efficient-multihead-attention/mha-implementations.ipynb
+++ b/ch03/02_bonus_efficient-multihead-attention/mha-implementations.ipynb
@@ -22,50 +22,6 @@
     "</table>"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "id": "1HABx0Hr3PDD",
-   "metadata": {
-    "id": "1HABx0Hr3PDD"
-   },
-   "source": [
-    "Uncomment and execute the following code cell to install the dependencies:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "id": "qPnVNAOxwy5s",
-   "metadata": {
-    "id": "qPnVNAOxwy5s"
-   },
-   "outputs": [],
-   "source": [
-    "# pip install -r https://raw.githubusercontent.com/rasbt/LLMs-from-scratch/main/requirements.txt"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "LYLcq3403Yq6",
-   "metadata": {
-    "id": "LYLcq3403Yq6"
-   },
-   "source": [
-    "Uncomment and execute the following code cell to install the PyTorch nightly dependency if you want to run the FlexAttention benchmarks (this is required because FlexAttention is not yet included in the latest PyTorch release):"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "id": "gAgYvxm_xVct",
-   "metadata": {
-    "id": "gAgYvxm_xVct"
-   },
-   "outputs": [],
-   "source": [
-    "# pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121 -U"
-   ]
-  },
   {
    "cell_type": "markdown",
    "id": "6f678e62-7bcb-4405-86ae-dce94f494303",
@@ -119,6 +75,28 @@
     "embeddings = torch.randn((batch_size, context_len, embed_dim), device=device)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "LYLcq3403Yq6",
+   "metadata": {
+    "id": "LYLcq3403Yq6"
+   },
+   "source": [
+    "- To run all the code in this notebook, please ensure you update to at least PyTorch 2.5 (FlexAttention is not included in earlier PyTorch releases)\n",
+    "If the code cell above shows a PyTorch version lower than 2.5, you can upgrade your PyTorch installation by uncommenting and running the following code cell (Please note that PyTorch 2.5 requires Python 3.9 or later)\n",
+    "- For more specific instructions and CUDA versions, please refer to the official installation guide at https://pytorch.org."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1db27f43-86f4-478f-89df-fbc2182a129b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# pip install --upgrade torch torchvision torchaudio"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "2f9bb1b6-a1e5-4e0a-884d-0f31b374a8d6",
@@ -964,16 +942,16 @@
     "## 9) Using PyTorch's FlexAttention\n",
     "\n",
     "- See [FlexAttention: The Flexibility of PyTorch with the Performance of FlashAttention](https://pytorch.org/blog/flexattention/) to learn more about FlexAttention\n",
-    "- This is currently only supported in PyTorch 2.5 (nightly), which you can install on a CPU machine via\n",
+    "- This is supported starting from PyTorch 2.5, which you can install on a CPU machine via\n",
     "\n",
     "    ```bash\n",
-    "    pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu -U\n",
+    "    pip install torch torchvision torchaudio\n",
     "    ```\n",
     "\n",
-    "- To install PyTorch nighly on a GPU machine, use the following (for more information, also see the installation menu on [pytorch.org](https://pytorch.org/))\n",
+    "- To install PyTorch on a GPU machine, use the following (for more information, also see the installation menu on [pytorch.org](https://pytorch.org/))\n",
     "\n",
     "    ```bash\n",
-    "    pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121 -U\n",
+    "    pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124\n",
     "    ```"
    ]
   },
@@ -2001,7 +1979,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.9"
+   "version": "3.11.4"
   }
  },
  "nbformat": 4,
diff --git a/setup/03_optional-docker-environment/.devcontainer/Dockerfile b/setup/03_optional-docker-environment/.devcontainer/Dockerfile
index 7e0f697..7919433 100644
--- a/setup/03_optional-docker-environment/.devcontainer/Dockerfile
+++ b/setup/03_optional-docker-environment/.devcontainer/Dockerfile
@@ -1,5 +1,7 @@
-FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
+# Install PyTorch 2.5 with CUDA 12.4
+FROM pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime
 
+# Install Ubuntu packages
 RUN apt-get update && \
     apt-get upgrade -y && \
     apt-get install -y rsync && \
@@ -7,6 +9,7 @@ RUN apt-get update && \
     apt-get install -y curl && \
     rm -rf /var/lib/apt/lists/*
 
+# Install Python packages
 COPY requirements.txt requirements.txt
-
+RUN pip install --upgrade pip
 RUN pip install --no-cache-dir -r requirements.txt