diff --git a/Dockerfile b/Dockerfile
index c84368a..4a55860 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -31,7 +31,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
 ADD src .  
 
 # Quick temporary updates
-RUN pip install git+https://github.com/runpod/runpod-python@main#egg=runpod --compile
+RUN pip install git+https://github.com/runpod/runpod-python@a1#egg=runpod --compile
 
 # Prepare the models inside the docker image
 ARG HUGGING_FACE_HUB_TOKEN=
@@ -55,7 +55,7 @@ ENV TRANSFORMERS_CACHE="/runpod-volume/huggingface-cache/hub"
 
 # Download the models
 RUN mkdir -p /model
-# RUN MODEL_NAME=$MODEL_NAME MODEL_REVISION=$MODEL_REVISION MODEL_BASE_PATH=$MODEL_BASE_PATH HUGGING_FACE_HUB_TOKEN=$HUGGING_FACE_HUB_TOKEN python -u /download_model.py
+RUN MODEL_NAME=$MODEL_NAME MODEL_REVISION=$MODEL_REVISION MODEL_BASE_PATH=$MODEL_BASE_PATH HUGGING_FACE_HUB_TOKEN=$HUGGING_FACE_HUB_TOKEN python -u /download_model.py
 
 # Start the handler
 CMD STREAMING=$STREAMING MODEL_NAME=$MODEL_NAME MODEL_BASE_PATH=$MODEL_BASE_PATH TOKENIZER=$TOKENIZER python -u /handler.py