diff --git a/irt_scripts/run_predict_results.sbatch b/irt_scripts/run_predict_results.sbatch
index 0172f3193..2835cd0eb 100644
--- a/irt_scripts/run_predict_results.sbatch
+++ b/irt_scripts/run_predict_results.sbatch
@@ -16,9 +16,12 @@ mkdir -p ${BASE_PATH}/experiments/predict_files/${model}
 
 echo "${model} ${task} ${config_no} ${model_path}"
 
+SHORT_MODEL_NAME="${MODEL_TYPE##*/}"
 
 python $JIANT_PATH/proj/main/runscript.py run   \
-    --ZZsrc ${BASE_PATH}/models/${model}/config.json  \
+    --hf_pretrained_model_name_or_path ${MODEL_TYPE} \
+    --model_path ${MODELS_DIR}/${SHORT_MODEL_NAME}/model/model.p \
+    --model_config_path ${MODELS_DIR}/${SHORT_MODEL_NAME}/model/config.json \
     --jiant_task_container_config_path ${BASE_PATH}/experiments/run_config_dir/taskmaster/${model}/${task}_${config_no}/${task}.json  \
     --model_load_mode all --model_path ${BASE_PATH}/experiments/output_dir/taskmaster_${model}_bestconfig/${task}/config_${config_no}/${model_path}  \
     --output_dir  ${BASE_PATH}/experiments/predict_files/${model}/${task}_config_${config_no}_${model_path}  \
diff --git a/irt_scripts/run_preprocess.sh b/irt_scripts/run_preprocess.sh
index d18db36d7..c12f6919e 100644
--- a/irt_scripts/run_preprocess.sh
+++ b/irt_scripts/run_preprocess.sh
@@ -24,26 +24,22 @@ function prepare_all_tasks() {
 
 
 function preprocess_task(){
+    # Full model name e.g. nyu-mll/roberta-base-100M-1, roberta-base
     MODEL_TYPE=$1
     TASK_NAME=$2
 
-    if [[ $MODEL_TYPE = nyu* ]]
-    then
-        MODEL_NAME="${MODEL_TYPE##*/}"
-    else
-        MODEL_NAME=${MODEL_TYPE}
-    fi
-    echo "$MODEL_NAME: ${TASK_NAME}, ${DATA_DIR}"
+    # Model name e.g. roberta-base-100M-1, roberta-base
+    SHORT_MODEL_NAME="${MODEL_TYPE##*/}"
+    echo "${SHORT_MODEL_NAME}: ${TASK_NAME}, ${DATA_DIR}"
 
     python ${JIANT_PATH}/proj/main/tokenize_and_cache.py \
-	    --task_config_path ${DATA_DIR}/configs/${TASK_NAME}_config.json \
-            --model_type ${MODEL_NAME} \
-            --model_tokenizer_path ${MODELS_DIR}/${MODEL_NAME}/tokenizer \
-            --phases train,val,test \
-            --max_seq_length 256 \
-	    --do_iter \
-            --smart_truncate \
-	    --output_dir ${CACHE_DIR}/${MODEL_NAME}/${TASK_NAME}
+        --task_config_path ${DATA_DIR}/configs/${TASK_NAME}_config.json \
+        --hf_pretrained_model_name_or_path ${MODEL_TYPE} \
+        --phases train,val,test \
+        --max_seq_length 256 \
+        --do_iter \
+        --smart_truncate \
+        --output_dir ${CACHE_DIR}/${SHORT_MODEL_NAME}/${TASK_NAME}
 }
 
 
diff --git a/irt_scripts/run_train_task.sbatch b/irt_scripts/run_train_task.sbatch
index 94efef571..abdaf0745 100644
--- a/irt_scripts/run_train_task.sbatch
+++ b/irt_scripts/run_train_task.sbatch
@@ -15,11 +15,14 @@ echo $RUN_CONFIG_DIR
 echo $SEED
 export PYTHONPATH=jiant/
 
+SHORT_MODEL_NAME="${MODEL_TYPE##*/}"
 
 python $JIANT_PATH/proj/main/runscript.py   \
     run_with_continue \
-       --ZZsrc ${MODELS_DIR}/${MODEL_TYPE}/config.json \
        --jiant_task_container_config_path ${RUN_CONFIG_DIR} \
+       --hf_pretrained_model_name_or_path ${MODEL_TYPE} \
+       --model_path ${MODELS_DIR}/${SHORT_MODEL_NAME}/model/model.p \
+       --model_config_path ${MODELS_DIR}/${SHORT_MODEL_NAME}/model/config.json \
        --model_load_mode from_transformers \
        --learning_rate $LR \
        --force_overwrite \