diff --git a/Dockerfile b/Dockerfile
index 71773da..5cf7c94 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -2,17 +2,18 @@ FROM ubuntu:latest
 RUN apt-get update -y
 RUN apt-get install -y python3 python3-pip
 
-RUN python3 -m pip install --user kaggle
-RUN python3 -m pip install --user pandas
-RUN python3 -m pip install --user scikit-learn
+RUN python3 -m pip install kaggle
+RUN python3 -m pip install pandas
+RUN python3 -m pip install scikit-learn
+RUN python3 -m pip install --upgrade tensorflow
 
 WORKDIR app
 
 ARG CUTOFF
 ENV CUTOFF=${CUTOFF}
 
-COPY main_docker.py ./
+COPY ml_training.py ./
 COPY heart_2020_cleaned.csv ./
 
 
-CMD ["python3", "./main_docker.py"]
+CMD ["python3", "./ml_training.py"]