Update for Dockerfile to change to cu128 to enable Blackwell GPU
This commit is contained in:
parent
7bd4cfceef
commit
3204232cd4
1 changed files with 6 additions and 3 deletions
|
|
@ -29,7 +29,8 @@ RUN python3 -m venv /app/venv
|
|||
ENV PATH="/app/venv/bin:$PATH"
|
||||
|
||||
# Install PyTorch with CUDA support and other dependencies
|
||||
RUN pip3 install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 && \
|
||||
# Using CUDA 12.8 for Blackwell sm_120 support
|
||||
RUN pip3 install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 && \
|
||||
pip3 install --no-cache-dir -r requirements.txt
|
||||
|
||||
# Copy project files
|
||||
|
|
@ -38,10 +39,12 @@ COPY --chown=appuser:appuser . .
|
|||
# Set environment variables
|
||||
ENV PYTHONUNBUFFERED=1 \
|
||||
PYTHONPATH=/app \
|
||||
USE_GPU=true
|
||||
USE_GPU=true \
|
||||
CUDA_LAUNCH_BLOCKING=1 \
|
||||
TORCH_USE_CUDA_DSA=1
|
||||
|
||||
# Expose the port
|
||||
EXPOSE 5005
|
||||
|
||||
# Run FastAPI server with uvicorn
|
||||
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "5005", "--workers", "1"]
|
||||
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "5005", "--workers", "1"]
|
||||
|
|
|
|||
Loading…
Reference in a new issue