diff --git a/job.pbs b/job.pbs
index 83faa1c..154dc9e 100644
--- a/job.pbs
+++ b/job.pbs
@@ -1,35 +1,40 @@
 #!/bin/bash
 
-#PBS -N neural-compression-test-run
+#PBS -N nc-cnn-enwik9-optuna
 #PBS -l gpus=1
-#PBS -l walltime=00:30:00
+#PBS -l walltime=08:00:00
+#PBS -l mem=60gb
 #PBS -m abe
 
 CACHE_DIR="${VSC_SCRATCH}/.cache"                # Directory to use as cache
 UV_DIR="${VSC_SCRATCH}/uv"                       # Directory to install packages
+VENV="${UV_DIR}/venv"
 
 DATA_DIR="${VSC_DATA}/datasets"
 RESULTS_DIR="${VSC_DATA}/neural-compression/$( date +%Y%m%d-%H%M-%S%N)-results"
 
+mkdir -p "${DATA_DIR}" "${RESULTS_DIR}" || true
+
 module purge
 module load PyTorch-bundle/2.1.2-foss-2023a-CUDA-12.1.1
 module load Optuna/3.5.0-foss-2023a
 module load matplotlib/2.2.5-foss-2023a-Python-2.7.18
 
-cd "${PBS_O_WORKDIR}"
+cd "${PBS_O_WORKDIR}" || exit
 
 UV_PYTHON_INSTALL_DIR="${UV_DIR}/python" UV_PYTHON_INSTALL_DIR="${UV_DIR}/python" \
   uv --cache-dir="${CACHE_DIR}/uv" \
-  venv "${UV_DIR}/venv" --clear
+  venv "${VENV}" --clear
 
-source "${UV_DIR}/venv/bin/activate"
+source "${VENV}/bin/activate"
 
 UV_PYTHON_INSTALL_DIR="${UV_DIR}/python" UV_PYTHON_INSTALL_DIR="${UV_DIR}/python" \
   uv --cache-dir="${CACHE_DIR}/uv" \
   sync --active --no-dev
 
-cd "${PBS_O_WORKDIR}/CNN-model"
+cd "${PBS_O_WORKDIR}" || exit
 
-python main_cnn.py train \
+python main.py train \
+  --method=optuna \
   --dataset=enwik9 --data-root="${DATA_DIR}" \
-  --model-path="${RESULTS_DIR}/model.pt"
+  --model=cnn --model-save-path="${RESULTS_DIR}/cnn-enwik9-optuna.pt"