fix: numpy

2024-09-12 14:19:35 -07:00
parent c576633fe0
commit 4c55f8ec71
1 changed files with 2 additions and 0 deletions
@@ -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