haritsahm commited on
Commit
7a40e74
·
1 Parent(s): a1f1417

Update dockerfile to use cuda image

Browse files
Files changed (1) hide show
  1. Dockerfile +2 -2
Dockerfile CHANGED
@@ -1,4 +1,4 @@
1
- from ubuntu:20.04
2
 
3
  ENV GRADIO_SERVER_NAME=0.0.0.0
4
 
@@ -13,7 +13,7 @@ RUN apt update && \
13
  rm -rf /var/lib/apt/lists
14
 
15
  RUN python3 -m pip install --no-cache-dir --upgrade pip && \
16
- python3 -m pip install --no-cache-dir --extra-index-url https://download.pytorch.org/whl/cpu \
17
  torch \
18
  torchvision \
19
  torchaudio
 
1
+ from nvidia/cuda:11.3.1-runtime-ubuntu20.04
2
 
3
  ENV GRADIO_SERVER_NAME=0.0.0.0
4
 
 
13
  rm -rf /var/lib/apt/lists
14
 
15
  RUN python3 -m pip install --no-cache-dir --upgrade pip && \
16
+ python3 -m pip install --no-cache-dir --extra-index-url https://download.pytorch.org/whl/cu113 \
17
  torch \
18
  torchvision \
19
  torchaudio