diff --git a/Dockerfile b/Dockerfile
index 0ed5be9..51f0853 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-FROM ubuntu:latest
+FROM ubuntu:20.04
 
 RUN apt-get update && apt-get install -y python3-pip && pip3 install setuptools && pip3 install numpy && pip3 install pandas && pip3 install wget && pip3 install scikit-learn && pip3 install matplotlib && rm -rf /var/lib/apt/lists/*
 RUN pip3 install torch torchvision torchaudio
diff --git a/dvc.yaml b/dvc.yaml
index 5fe1865..ae96f4e 100644
--- a/dvc.yaml
+++ b/dvc.yaml
@@ -1,9 +1,8 @@
 stages:
   download_and_split:
-    cmd: python3 split_10.py
+    cmd: python3 create.py
     deps:
     - healthcare-dataset-stroke-data.csv
-    - create.py
     outs:
     - data_train.csv
     - data_test.csv