Skip to content

Commit

Permalink
onnx gpu
Browse files Browse the repository at this point in the history
  • Loading branch information
csokun committed Jul 31, 2024
1 parent 7c0e8b0 commit 12d922a
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ RUN apt-get install -y libcudnn8
RUN apt-get install libgl1-mesa-glx -y

RUN --mount=type=cache,target=/root/.cache \
# pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu121
pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu124

WORKDIR /app
Expand All @@ -26,6 +25,7 @@ RUN --mount=type=cache,target=/root/.cache \
pip install -r requirements_docker.txt -r requirements.txt

RUN --mount=type=cache,target=/root/.cache \
pip uninstall -y onnxruntime onnxruntime-gpu && \
pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/

COPY . .
Expand Down

0 comments on commit 12d922a

Please sign in to comment.