diff --git a/.runpod/Dockerfile b/.runpod/Dockerfile
index d688d2b31..dfd8e1c16 100644
--- a/.runpod/Dockerfile
+++ b/.runpod/Dockerfile
@@ -1,4 +1,4 @@
-FROM axolotlai/axolotl-cloud:main-py3.11-cu124-2.6.0
+FROM runpod/pytorch:3.10-2.0.0-117
 
 COPY .runpod/requirements.txt /requirements.txt
 RUN --mount=type=cache,target=/root/.cache/pip \
diff --git a/.runpod/requirements.txt b/.runpod/requirements.txt
index a4a6fe21a..4f9e20da7 100644
--- a/.runpod/requirements.txt
+++ b/.runpod/requirements.txt
@@ -5,6 +5,11 @@
 # git+https://github.com/runpod/runpod-python.git
 # To learn more, see https://pip.pypa.io/en/stable/reference/requirements-file-format/
 runpod~=1.7.0
+huggingface_hub
 typing-extensions
 pydantic
 pydantic-settings
+hf-transfer
+setuptools
+numpy==2.0.0
+axolotl[flash-attn,deepspeed]