diff --git a/Dockerfile b/Dockerfile
index e0e23e8..73ac242 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,5 +1,5 @@
 # Use the official Python 3.8 image as the base image
-FROM pytorch/pytorch:2.3.1-cuda11.8-cudnn8-devel
+FROM pytorch/pytorch:2.3.1-cuda11.8-cudnn8-runtime
 
 # Set the working directory inside the container
 WORKDIR /app