diff --git a/Dockerfile b/Dockerfile
index 4b53120..5dd36a8 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -10,12 +10,12 @@ ENV DEBIAN_FRONTEND=noninteractive \
     USE_CUDA=1 USE_DISTRIBUTED=0 BUILD_TEST=0 TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;9.0"
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
-      python3 python3-pip python3-distutils git cmake ninja-build \
+      python3 python3-dev python3-pip python3-distutils git cmake ninja-build \
       libopenblas-dev libopenmpi-dev \
       libnccl2=2.22.3-1+cuda12.6 \
       libnccl-dev=2.22.3-1+cuda12.6 \
       libjpeg-dev libpng-dev ca-certificates && \
-    python3 -m pip install --no-cache-dir --upgrade pip wheel setuptools sympy pyyaml typing-extensions
+    python3 -m pip install --no-cache-dir --upgrade pip wheel setuptools sympy pyyaml typing-extensions numpy
 
 WORKDIR /opt
 RUN git clone --recursive -b v2.7.1 https://github.com/pytorch/pytorch.git