fix: gpu arch

This commit is contained in:
2024-09-10 15:05:41 -07:00
parent 476aa7391c
commit af74c897ad
+2 -1
View File
@@ -5,10 +5,11 @@ FROM wtongze/rocm-base:${ROCM_VERSION}
ARG TENSORFLOW_BRANCH="r2.11-rocm-enhanced" ARG TENSORFLOW_BRANCH="r2.11-rocm-enhanced"
ARG PYTHON_VERSION="3.11" ARG PYTHON_VERSION="3.11"
ARG NUMPY_VERSION="1.22.4" ARG NUMPY_VERSION="1.22.4"
ARG HIP_ARCH="gfx803;gfx900"
ARG DEBIAN_FRONTEND="noninteractive" ARG DEBIAN_FRONTEND="noninteractive"
# Add GPU target # Add GPU target
RUN bash -c 'echo -e "gfx803\ngfx900" | sudo tee /opt/rocm/bin/target.lst' RUN echo "${HIP_ARCH}" | sed 's/;/\n/g' | sudo tee /opt/rocm/bin/target.lst
# Install bazelisk # Install bazelisk
RUN sudo wget -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.21.0/bazelisk-linux-amd64 && sudo chmod 755 /usr/local/bin/bazel RUN sudo wget -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.21.0/bazelisk-linux-amd64 && sudo chmod 755 /usr/local/bin/bazel