diff --git a/Dockerfile b/Dockerfile
index edd189c..9f9fc51 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -28,5 +28,5 @@ RUN chmod +x ./train.py
 ARG epochs=5
 RUN echo $epochs
 
-CMD python3 ./train.py epochs
+CMD python3 ./train.py $epochs
 #CMD python3 ./test.py
\ No newline at end of file