diff --git a/aimodel/slurm-TEST-deeplabv3p-rainfall.job b/aimodel/slurm-TEST-deeplabv3p-rainfall.job
index 6e0e792..89d1eea 100755
--- a/aimodel/slurm-TEST-deeplabv3p-rainfall.job
+++ b/aimodel/slurm-TEST-deeplabv3p-rainfall.job
@@ -59,7 +59,7 @@ DIR_OUTPUT="output/$(date -u --rfc-3339=date)_${CODE}";
 echo -e ">>> Additional args: ${ARGS}";
 
 export PATH=$HOME/software/bin:$PATH;
-export IMAGE_SIZE BATCH_SIZE DIR_RAINFALLWATER PATH_HEIGHTMAP PATH_COLOURMAP STEPS_PER_EPOCH DIR_OUTPUT PATH_CHECKPOINT EPOCHS;
+export IMAGE_SIZE BATCH_SIZE DIR_RAINFALLWATER PATH_HEIGHTMAP PATH_COLOURMAP STEPS_PER_EPOCH DIR_OUTPUT PATH_CHECKPOINT EPOCHS PREDICT_COUNT;
 
 echo ">>> Installing requirements";
 conda run -n py38 pip install -q -r requirements.txt;