diff --git a/pytorch.dockerfile b/pytorch.dockerfile
index 91a9a37..7e38458 100644
--- a/pytorch.dockerfile
+++ b/pytorch.dockerfile
@@ -6,6 +6,7 @@ ARG ROCM_VERSION="5.7.3"
 ARG PYTORCH_BRANCH="v2.2.2"
 ARG TORCHVISION_BRANCH="v0.17.2"
 ARG PYTHON_VERSION="3.11"
+ARG NUMPY_VERSION="1.26.4"
 ARG HIP_ARCH="gfx803;gfx900"
 ARG DEBIAN_FRONTEND="noninteractive"
 
@@ -29,6 +30,7 @@ EOF
 
 RUN python${PYTHON_VERSION} -m pip install wheel
 RUN python${PYTHON_VERSION} -m pip install -r requirements.txt
+RUN python${PYTHON_VERSION} -m pip install numpy==${NUMPY_VERSION}
 
 RUN <<EOF
 sudo apt-get update