Gemma notebook working on runpod

Dockerfile

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.8.1-cudnn-runtime-ubuntu24.04 AS base
+FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 as base
 RUN apt-get update && \
     apt-get upgrade -y && \
     apt install -y \
@@ -6,13 +6,13 @@ RUN apt-get update && \
     git \
     gnupg \
     jq \
-    python3-pip \
+    python3-dev \
     python3-venv \
     tmux \
     vim \
     wget
 
-RUN wget -q -O - https://dl-ssl.google.com/linux/linux_signing_key.pub | apt-key add -
+RUN wget -q -O - https://dl-ssl.google.com/linux/linux_signing_key.pub | gpg --dearmor -o /etc/apt/trusted.gpg.d/google.gpg
 RUN echo "deb http://dl.google.com/linux/chrome/deb/ stable main" >> /etc/apt/sources.list.d/google-chrome.list
 RUN apt-get update && \
     apt-get install -y \
@@ -63,6 +63,16 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     /root/venvs/sia/bin/pip install -r *.egg-info/requires.txt
 RUN rm -rf *.egg-info/
 
+# Notebook setup
+FROM base AS notebook-env
+RUN python3 -m venv /root/venvs/notebook --system-site-packages
+RUN --mount=type=cache,target=/root/.cache/pip \
+    /root/venvs/notebook/bin/pip install \
+        jupyter \
+        ipykernel \
+        ipywidgets
+RUN /root/venvs/notebook/bin/ipython kernel install --name=notebook
+
 # Web frontend build
 FROM base
 # AS web-build
@@ -88,6 +98,7 @@ FROM base
 COPY --from=itb-env /root/venvs/itb /root/venvs/itb
 COPY --from=train-env /root/venvs/train /root/venvs/train
 COPY --from=sia-env /root/venvs/sia /root/venvs/sia
+COPY --from=notebook-env /root/venvs/notebook /root/venvs/notebook
 COPY --from=web-build /app/dist /root/static/
 
 RUN echo 'source /root/sia/scripts/add_venvs_to_path.sh' >> /root/.bashrc