Option to override PYTORCH_ROCM_ARCH inherited from the base image (#321

)
ROCm · Dec 11, 2024 · a1aaa74 · a1aaa74
1 parent 2bec3da
commit a1aaa74
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/Dockerfile.rocm b/Dockerfile.rocm
@@ -7,6 +7,9 @@ ARG BASE_IMAGE=rocm/vllm-dev:base_ubuntu22.04_py3.12_ROCm6.3_hipblaslt0.12_torch
 
 FROM ${BASE_IMAGE} AS base
 
+ARG ARG_PYTORCH_ROCM_ARCH
+ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
+
 # Install some basic utilities
 RUN apt-get update -q -y && apt-get install -q -y \
     sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev