feat: simplify rocm 4.5+

2024-09-11 21:01:57 -07:00
parent bd42ce0763
commit 2b2d578ab8
3 changed files with 8 additions and 16 deletions
@@ -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