diff --git a/dvc.yaml b/dvc.yaml
index cfce856..9c05247 100644
--- a/dvc.yaml
+++ b/dvc.yaml
@@ -1,6 +1,6 @@
 stages:
   prepare_datasets:
-    cmd: python prepare_datasets.py X_train.csv X_test.csv y_train.csv y_test.csv
+    cmd: python prepare_datasets.py
     deps:
     - data/bodyPerformance.csv
     - prepare_datasets.py
@@ -10,7 +10,7 @@ stages:
     - y_test.csv
     - y_train.csv
   train:
-    cmd: python train.py model.pth
+    cmd: python train.py
     deps:
     - X_train.csv
     - train.py