浏览代码

Update PyTorch wheel source to CUDA 11.8

Modified the Dockerfile to install PyTorch, torchvision, and torchaudio from a CUDA 11.8 specific wheel URL. This ensures compatibility with the CUDA version in our environment and potentially improves performance and stability for GPU-accelerated operations.
Joseph Young 1 年之前
父节点
当前提交
2588da0e27
共有 1 个文件被更改,包括 1 次插入1 次删除
  1. 1 1
      Dockerfile-cuda

+ 1 - 1
Dockerfile-cuda

@@ -86,7 +86,7 @@ RUN apk update && \
     apk del /var/cache/apk/*.tbz2
 
 COPY ./backend/requirements.txt ./requirements.txt
-RUN pip3 install torch torchvision torchaudio --no-cache-dir && \
+RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 --no-cache-dir && \
     pip3 install -r requirements.txt --no-cache-dir
 
 # copy built frontend files