diff --git a/Dockerfile b/Dockerfile
index 5812ceb..c95d4a3 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,5 +1,5 @@
 # Stage 1: Build stage
-FROM pytorch:slim AS builder
+FROM python:slim AS builder
 
 WORKDIR /app