diff --git a/pytorch.dockerfile b/pytorch.dockerfile
index 8e1b95b..8889e55 100644
--- a/pytorch.dockerfile
+++ b/pytorch.dockerfile
@@ -49,7 +49,7 @@ ENV PYTORCH_ROCM_ARCH=${HIP_ARCH}
 ENV PYTORCH_BUILD_VERSION=${PYTORCH_BUILD_VERSION}
 ENV PYTORCH_BUILD_NUMBER=0
 RUN python${PYTHON_VERSION} tools/amd_build/build_amd.py && \
-USE_ROCM=ON USE_FLASH_ATTENTION=OFF USE_MEM_EFF_ATTENTION=OFF python${PYTHON_VERSION} setup.py bdist_wheel
+USE_ROCM=ON USE_NUMPY=1 USE_FLASH_ATTENTION=OFF USE_MEM_EFF_ATTENTION=OFF python${PYTHON_VERSION} setup.py bdist_wheel
 
 # https://github.com/pytorch/pytorch/issues/125230
 # https://github.com/AngryLoki/pytorch/commit/38d6f115c65b67946de6765165ece0288c8416d4