diff --git a/phases/00-setup-and-tooling/07-docker-for-ai/code/Dockerfile b/phases/00-setup-and-tooling/07-docker-for-ai/code/Dockerfile
index 82958af906..4419a9d1c2 100644
--- a/phases/00-setup-and-tooling/07-docker-for-ai/code/Dockerfile
+++ b/phases/00-setup-and-tooling/07-docker-for-ai/code/Dockerfile
@@ -3,27 +3,33 @@ FROM nvidia/cuda:12.4.1-devel-ubuntu22.04
 ENV DEBIAN_FRONTEND=noninteractive
 ENV PYTHONUNBUFFERED=1
 
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    python3.12 \
-    python3.12-venv \
-    python3.12-dev \
-    python3-pip \
-    git \
-    curl \
-    build-essential \
+RUN apt-get update && apt-get install -y software-properties-common && \
+    add-apt-repository ppa:deadsnakes/ppa && \
+    apt-get update && apt-get install -y --no-install-recommends \
+        python3.12 \
+        python3.12-venv \
+        python3.12-dev \
+        git \
+        curl \
+        build-essential \
     && rm -rf /var/lib/apt/lists/*
 
-RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.12 1
+RUN curl -sSf https://bootstrap.pypa.io/get-pip.py -o /tmp/get-pip.py && \
+    python3.12 /tmp/get-pip.py && \
+    rm -f /tmp/get-pip.py
 
-RUN python -m pip install --no-cache-dir --upgrade pip setuptools wheel
+RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.12 1 && \
+    update-alternatives --install /usr/bin/pip pip /usr/local/bin/pip3.12 1
 
-RUN python -m pip install --no-cache-dir \
-    torch==2.3.1 \
-    torchvision==0.18.1 \
-    torchaudio==2.3.1 \
+RUN pip install --no-cache-dir --upgrade pip setuptools wheel
+
+RUN pip install --no-cache-dir \
+    torch==2.6.0+cu124 \
+    torchvision==0.21.0+cu124 \
+    torchaudio==2.6.0+cu124 \
     --index-url https://download.pytorch.org/whl/cu124
 
-RUN python -m pip install --no-cache-dir \
+RUN pip install --no-cache-dir \
     numpy \
     pandas \
     scikit-learn \