diff --git a/Dockerfile.base b/Dockerfile.base
index 3700920..d1a1eec 100644
--- a/Dockerfile.base
+++ b/Dockerfile.base
@@ -1,5 +1,5 @@
 # Base image with all ML dependencies
-FROM python:3.13-bookworm
+FROM pytorch/pytorch:2.8.0-cuda12.6-cudnn9-runtime
 
 # Install system dependencies
 RUN apt update && apt install -y libgl1 && rm -rf /var/lib/apt/lists/*