소스 검색

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