blanchon commited on
Commit
7ab95fc
·
1 Parent(s): ebc839f
Files changed (1) hide show
  1. Dockerfile +2 -2
Dockerfile CHANGED
@@ -2,7 +2,7 @@
2
  FROM ghcr.io/astral-sh/uv:0.2.12 as uv
3
 
4
  # Main application image
5
- FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04
6
 
7
  # Set home to the user's home directory
8
  ENV DEBIAN_FRONTEND=noninteractive \
@@ -58,7 +58,7 @@ COPY --chown=user requirements.txt .
58
 
59
  # Install Python packages with uv caching
60
  RUN --mount=type=cache,target=$UV_CACHE_DIR,uid=1000,gid=1000 \
61
- /uv pip install torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/cu118
62
 
63
  # Install build dependencies for flash-attn
64
  RUN --mount=type=cache,target=$UV_CACHE_DIR,uid=1000,gid=1000 \
 
2
  FROM ghcr.io/astral-sh/uv:0.2.12 as uv
3
 
4
  # Main application image
5
+ FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04
6
 
7
  # Set home to the user's home directory
8
  ENV DEBIAN_FRONTEND=noninteractive \
 
58
 
59
  # Install Python packages with uv caching
60
  RUN --mount=type=cache,target=$UV_CACHE_DIR,uid=1000,gid=1000 \
61
+ /uv pip install torch torchvision
62
 
63
  # Install build dependencies for flash-attn
64
  RUN --mount=type=cache,target=$UV_CACHE_DIR,uid=1000,gid=1000 \