diff --git a/Dockerfile b/Dockerfile
index 1dce7b4..0760252 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -11,5 +11,7 @@ WORKDIR /app
 
 COPY ./download.sh ./
 COPY ./calculate.sh ./
+COPY ./train.py ./
+COPY ./test.py ./
 
-CMD bash
\ No newline at end of file
+CMD bash download.sh && bash calculate.sh && python train.py && python test.py
\ No newline at end of file
diff --git a/download.sh b/download.sh
index 08c2121..c1cb981 100644
--- a/download.sh
+++ b/download.sh
@@ -2,4 +2,5 @@ pip install kaggle
 kaggle datasets download -d fanconic/skin-cancer-malignant-vs-benign
 unzip skin-cancer-malignant-vs-benign.zip 
 rm skin-cancer-malignant-vs-benign.zip
+rm -r data
 
diff --git a/requirements.txt b/requirements.txt
index 3ea24b4..c39b986 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,5 @@
 matplotlib==3.8.3
 numpy==1.26.4
+torchvision==0.17.1
+torch==2.2.1