UPDATE

Dockerfile

@@ -19,14 +19,18 @@ ENV PATH=/root/.cargo/bin:$PATH
 RUN pnpm install
 RUN pnpm run build --mode target-rwkv-hf-space
 
-FROM nvidia/cuda:12.4.0-runtime-ubuntu22.04 AS Backend
+FROM nvidia/cuda:12.4.0-devel-ubuntu22.04 AS Backend
 
 RUN <<EOF
 apt update
 apt install --no-install-recommends -y \
     build-essential \
     git \
-    cuda-toolkit-12-4
+    cuda-nvcc-12-4 \
+    cuda-cudart-dev-12-4 \
+    python3-dev \
+    python3-pip \
+    libpython3.10-dev
 apt clean && rm -rf /var/lib/apt/lists/*
 EOF
 
@@ -39,11 +43,10 @@ RUN useradd -m -u 1000 user
 USER user
 
 ENV HOME=/home/user \
-    PATH=/home/user/.local/bin:$PATH \
-    PATH=/usr/local/cuda/bin:$PATH \
-    CXX=g++ \
-    LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
-
+    PATH=/usr/local/cuda/bin:/home/user/.local/bin:$PATH \
+    LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}" \
+    CXX=/usr/bin/g++ \
+    TORCH_CUDA_ARCH_LIST="7.5"  
 WORKDIR $HOME/app
 
 COPY --chown=user . $HOME/app