From e0cdcda4ba7990368d615dbbb664e1f97fa0efc6 Mon Sep 17 00:00:00 2001
From: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Date: Wed, 11 Dec 2024 13:29:50 -0600
Subject: [PATCH] Option to override PYTORCH_ROCM_ARCH inherited from the base
 image

---
 Dockerfile.rocm | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/Dockerfile.rocm b/Dockerfile.rocm
index b665bbd42709b..1f44958f5d1f1 100644
--- 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