diff --git a/Dockerfile.gpu b/Dockerfile.gpu
index 3ef3835..80c4912 100644
--- a/Dockerfile.gpu
+++ b/Dockerfile.gpu
@@ -29,7 +29,7 @@ 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 && \
+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
@@ -44,4 +44,4 @@ ENV PYTHONUNBUFFERED=1 \
 EXPOSE 5005
 
 # Run FastAPI server with uvicorn
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "5005", "--workers", "1"]
\ No newline at end of file
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "5005", "--workers", "1"]