diff --git a/stable-diffusion/setup_runpod_lora_learn.sh b/stable-diffusion/setup_runpod_lora_learn.sh
index 54c636d..6f713cf 100644
--- a/stable-diffusion/setup_runpod_lora_learn.sh
+++ b/stable-diffusion/setup_runpod_lora_learn.sh
@@ -20,6 +20,9 @@ pip install --upgrade \
 echo "Installing xformers"
 
 pip install --no-deps xformers==0.0.17
+
+echo "Installing TensorRT"
+
 pip install tensorrt
 
 echo "Cloning the project..."
@@ -33,6 +36,8 @@ TEMP_REQUIREMENTS_FILE="${DIR}/requirements_tmp_for_setup.txt"
 awk -v dir="$DIR" '/#.*kohya_ss.*library/{print; getline; sub(/^\.$/, dir)}1' "$DIR/requirements.txt" >"$TEMP_REQUIREMENTS_FILE"
 python -m pip install --use-pep517 --upgrade -r "$TEMP_REQUIREMENTS_FILE"
 
+echo "Installing accelerate config"
+
 mkdir -p "$HOME/.cache/huggingface/accelerate"
 cp "${DIR}/config_files/accelerate/default_config.yaml"\
     "$HOME/.cache/huggingface/accelerate/default_config.yaml"
@@ -45,5 +50,14 @@ pushd /workspace
 curl -L -O -J https://civitai.com/api/download/models/29792
 popd
 
+echo "Creating directories"
+
+mkdir -p /workspace/images/train/img \
+    /workspace/images/train/model \
+    /workspace/images/train/log
+
+echo "Cleanup"
+rm -r ~/.cache/huggingface ~/.cache/pip*
+
 print "Setup complete."
 exit 0