diff --git a/.gitignore b/.gitignore
index 2f9182e..c1daf76 100644
--- a/.gitignore
+++ b/.gitignore
@@ -70,6 +70,7 @@ fastlane/test_output
 Package.resolved
 .swiftpm/
 Public/images/
+Public/classifier.version
+Public/classifier.mlmodel
+Public/caps.json
 Training/backup/
-Training/classifier.version
-
diff --git a/Training/train.sh b/Training/train.sh
index 4b0b430..75bd0a5 100755
--- a/Training/train.sh
+++ b/Training/train.sh
@@ -30,8 +30,8 @@
 WORK_DIR="${HOME}/Projects/Caps/Caps-Server/Training"
 BACKUP_DIR="./backup"
 IMAGE_DIR="../Public/images"
-VERSION_FILE="./classifier.version"
-MODEL_FILE="./classifier.mlmodel"
+VERSION_FILE="../Public/classifier.version"
+MODEL_FILE="../Public/classifier.mlmodel"
 
 TRAINING_ITERATIONS="17"
 SSH_PORT="5432"
@@ -49,20 +49,6 @@ if [ $retVal -ne 0 ]; then
     return $retVal
 fi
 
-echo "[INFO] Getting classifier version from server..."
-scp -P $SSH_PORT ${SERVER}:/${SERVER_ROOT_PATH}/classifier.version .
-
-retVal=$?
-if [ $retVal -ne 0 ]; then
-    echo '[ERROR] Failed to get classifier version'
-    return $retVal
-fi
-
-# Read classifier version from file
-OLD_VERSION=$(< $VERSION_FILE)
-NEW_VERSION=$(($OLD_VERSION + 1))
-echo "[INFO] Creating classifier version ${NEW_VERSION}" 
-
 echo "[INFO] Ensuring permissions for images on server..."
 ssh -p $SSH_PORT ${SERVER} "sudo chmod -R 755 ${SERVER_ROOT_PATH}/images"
 
@@ -81,7 +67,28 @@ if [ $retVal -ne 0 ]; then
     return $retVal
 fi
 
-echo "[INFO] Training the model..."
+echo "[INFO] Getting classifier version from server..."
+scp -P $SSH_PORT ${SERVER}:/${SERVER_ROOT_PATH}/classifier.version $VERSION_FILE
+
+retVal=$?
+if [ $retVal -ne 0 ]; then
+    echo '[ERROR] Failed to get classifier version'
+    return $retVal
+fi
+
+# Read classifier version from file
+OLD_VERSION=$(< $VERSION_FILE)
+NEW_VERSION=$(($OLD_VERSION + 1))
+
+echo "[INFO] Backing up model ${OLD_VERSION}..."
+mv $MODEL_FILE "${BACKUP_DIR}/classifier${OLD_VERSION}.mlmodel"
+
+retVal=$?
+if [ $retVal -ne 0 ]; then
+    echo '[WARNING] Failed to back up old model'
+fi
+
+echo "[INFO] Training model ${NEW_VERSION} ..."
 swift train.swift $IMAGE_DIR $TRAINING_ITERATIONS $MODEL_FILE
 
 retVal=$?
@@ -102,7 +109,7 @@ if [ $retVal -ne 0 ]; then
     return $retVal
 fi
 
-echo "[INFO] Moving files into public directory..."
+echo "[INFO] Moving server files into public directory..."
 ssh -p ${SSH_PORT} ${SERVER} "sudo mv /home/pi/classifier.* ${SERVER_ROOT_PATH}"
 
 retVal=$?
@@ -111,7 +118,7 @@ if [ $retVal -ne 0 ]; then
     return $retVal
 fi
 
-echo "[INFO] Updating permissions..."
+echo "[INFO] Updating server permissions..."
 ssh -p ${SSH_PORT} ${SERVER} "sudo chown -R www-data\:www-data ${SERVER_ROOT_PATH}"
 
 retVal=$?
@@ -120,16 +127,4 @@ if [ $retVal -ne 0 ]; then
     return $retVal
 fi
 
-echo "[INFO] Backing up model..."
-mv $MODEL_FILE "${BACKUP_DIR}/classifier${NEW_VERSION}.mlmodel"
-
-retVal=$?
-if [ $retVal -ne 0 ]; then
-    echo '[WARNING] Failed to back up model'
-    rm $MODEL_FILE
-fi
-
-echo "[INFO] Cleaning up..."
-rm $VERSION_FILE
-
 echo "[INFO] Process finished"