diff --git a/aimodel/slurm-train.job b/aimodel/slurm-train.job
index f89e769..86368a7 100755
--- a/aimodel/slurm-train.job
+++ b/aimodel/slurm-train.job
@@ -63,7 +63,7 @@ echo ">>> Installing requirements";
 conda run -n py38 pip install -r requirements.txt;
 echo ">>> Training model";
 #shellcheck disable=SC2086
-/usr/bin/env time -v conda run -n py38 src/index.py pretrain --input "${INPUT}" --output "${dir_output}" ${ARGS};
+/usr/bin/env time -v conda run -n py38 src/index.py train --input "${INPUT}" --output "${dir_output}" ${ARGS};
 src/index.py train --input "${INPUT}" -o "${dir_output}"
 # Recommended extra args: --water-size 94 --batch-size 48 --arch convnext_i_xtiny