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