diff --git a/dvc.lock b/dvc.lock
new file mode 100644
index 0000000..01d79c5
--- /dev/null
+++ b/dvc.lock
@@ -0,0 +1,4 @@
+schema: '2.0'
+stages:
+  prepare:
+    cmd: download_dataset.sh
diff --git a/dvc.yaml b/dvc.yaml
index f92dd9a..640bd36 100644
--- a/dvc.yaml
+++ b/dvc.yaml
@@ -1,6 +1,6 @@
 stages:
   prepare:
-    cmd: ./download_dataset.sh
+    cmd: download_dataset.sh
     
   train:
     cmd: python3 train_model.py
@@ -11,4 +11,5 @@ stages:
     plots:
       - trend.png
     params:
-      - model_params.yaml
\ No newline at end of file
+      - batch_size
+      - epochs
\ No newline at end of file
diff --git a/model_params.yaml b/model_params.yaml
deleted file mode 100644
index cbf4e69..0000000
--- a/model_params.yaml
+++ /dev/null
@@ -1,3 +0,0 @@
-train:
-  batch_size: 64
-  epochs: 5
\ No newline at end of file
diff --git a/params.yaml b/params.yaml
new file mode 100644
index 0000000..795d3b6
--- /dev/null
+++ b/params.yaml
@@ -0,0 +1,2 @@
+batch_size: 64
+epochs: 5
\ No newline at end of file