feat: simplify rocm 4.5+

base.dockerfile

@@ -50,6 +50,14 @@ apt-get clean
 rm -rf /var/lib/apt/lists/*
 EOF
 
+# For ROCm 4.5+
+RUN <<EOF
+sudo apt-get update
+sudo apt-get install -y rocm-ml-sdk rocm-hip-sdk || true
+sudo apt-get clean
+sudo rm -rf /var/lib/apt/lists/*
+EOF
+
 # Create non-root user
 RUN <<EOF
 echo "%sudo ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/sudo-nopasswd

pytorch.dockerfile

@@ -29,14 +29,6 @@ EOF
 
 RUN python${PYTHON_VERSION} -m pip install typing-extensions wheel
 
-# For ROCm 4.5+
-RUN <<EOF
-sudo apt-get update
-sudo apt-get install -y rocm-ml-sdk rocm-hip-sdk
-sudo apt-get clean
-sudo rm -rf /var/lib/apt/lists/*
-EOF
-
 RUN <<EOF
 sudo apt-get update
 sudo apt-get install -y ninja-build

tensorflow.dockerfile

@@ -18,14 +18,6 @@ RUN sudo wget -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/rel
 WORKDIR /home/rocm
 RUN git clone --depth 1 https://github.com/ROCm/tensorflow-upstream.git -b ${TENSORFLOW_BRANCH} tensorflow
 
-# For ROCm 4.5+
-RUN <<EOF
-sudo apt-get update
-sudo apt-get install -y rocm-ml-sdk rocm-hip-sdk
-sudo apt-get clean
-sudo rm -rf /var/lib/apt/lists/*
-EOF
-
 RUN <<EOF
 sudo apt-get update
 sudo apt install -y software-properties-common