fix rocm torch version

ai-dock · Oct 15, 2023 · bb8bb1e · bb8bb1e
1 parent ede3e37
commit bb8bb1e
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/amd.sh b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/amd.sh
@@ -8,6 +8,12 @@ main() {
 }
 
 install_comfyui() {
+    # Mamba export does not include pip packages.
+    # We need to get torch again - todo find a better way?
+    micromamba -n comfyui run pip install \
+        --no-cache-dir \
+        --index-url https://download.pytorch.org/whl/rocm${ROCM_VERSION} \
+        torch==${PYTORCH_VERSION} torchvision torchaudio
     /opt/ai-dock/bin/update-comfyui.sh
 }