diff --git a/Dockerfile b/Dockerfile
index 11c74af..3bef3a2 100755
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,5 +1,5 @@
 # Stage 1: Builder
-FROM nvidia/cuda:12.6.3-base-ubuntu22.04 AS builder
+FROM nvidia/cuda:12.3.2-cudnn9-runtime-ubuntu22.04 AS builder
 
 WORKDIR /subgen
 
@@ -22,7 +22,7 @@ RUN pip install --no-cache-dir -r requirements.txt
 COPY . .
 
 # Stage 2: Runtime
-FROM nvidia/cuda:12.6.3-base-ubuntu22.04
+FROM nvidia/cuda:12.3.2-cudnn9-runtime-ubuntu22.04
 
 WORKDIR /subgen