diff --git a/Dockerfile b/Dockerfile
index 3937764..0b7a570 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,9 +1,12 @@
-# Base image
-FROM runpod/base:0.4.2-cuda12.1.0
+# Base image - Set default to CUDA 11.8.0
+ARG CUDA_VERSION=11.8.0
 
-ENV HF_DATASETS_CACHE="/runpod-volume/huggingface-cache/datasets"
-ENV HUGGINGFACE_HUB_CACHE="/runpod-volume/huggingface-cache/hub"
-ENV TRANSFORMERS_CACHE="/runpod-volume/huggingface-cache/hub"
+# Use different base images based on CUDA_VERSION argument
+FROM runpod/base:0.4.2-cuda${CUDA_VERSION} as builder
+
+ENV HF_DATASETS_CACHE="/runpod-volume/huggingface-cache/datasets" \
+    HUGGINGFACE_HUB_CACHE="/runpod-volume/huggingface-cache/hub" \
+    TRANSFORMERS_CACHE="/runpod-volume/huggingface-cache/hub"
 
 # Install Python dependencies (Worker Template)
 COPY builder/requirements.txt /requirements.txt
@@ -12,22 +15,26 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     python3.11 -m pip install --upgrade -r /requirements.txt --no-cache-dir && \
     rm /requirements.txt
 
+# Install specific packages based on CUDA version
+RUN if [ "$CUDA_VERSION" = "12.1.0" ]; then \
+        python3.11 -m pip install vllm==0.2.3; \
+    else \
+        python3.11 -m pip install vllm-0.2.3+cu118-cp311-cp311-manylinux1_x86_64.whl; \
+    fi
+
+
 # Add source files
 ADD src .
 
 ARG MODEL_NAME=""
-ENV MODEL_NAME=$MODEL_NAME
 ARG MODEL_BASE_PATH=""
-ENV MODEL_BASE_PATH=$MODEL_BASE_PATH
-
-# Set the environment variables conditionally
-RUN if [ -n "$MODEL_NAME" ]; then export MODEL_NAME=$MODEL_NAME; fi
-RUN if [ -n "$MODEL_BASE_PATH" ]; then export MODEL_BASE_PATH=$MODEL_BASE_PATH; fi
 
 # Conditionally run download_model.py
 RUN if [ -n "$MODEL_NAME" ] && [ -n "$MODEL_BASE_PATH" ]; then \
         python3.11 /download_model.py --model $MODEL_NAME --download_dir $MODEL_BASE_PATH; \
+        export MODEL_NAME=$MODEL_NAME; \
+        export MODEL_BASE_PATH=$MODEL_BASE_PATH; \
     fi
 
 # Start the handler
-CMD python3.11 /handler.py
+CMD ["python3.11", "/handler.py"]
\ No newline at end of file
diff --git a/builder/requirements.txt b/builder/requirements.txt
index f307ae9..d4ea3cb 100644
--- a/builder/requirements.txt
+++ b/builder/requirements.txt
@@ -1,4 +1,3 @@
 hf_transfer
 runpod==1.4.0
 huggingface-hub==0.19.4
-vllm==0.2.3
diff --git a/src/download_model.py b/src/download_model.py
index d0a0001..afb419e 100644
--- a/src/download_model.py
+++ b/src/download_model.py
@@ -1,4 +1,5 @@
 import argparse
+import os
 from vllm.model_executor.weight_utils import prepare_hf_model_weights
 
 if __name__ == "__main__":
@@ -10,6 +11,9 @@ if __name__ == "__main__":
     if not args.model or not args.download_dir:
         raise ValueError("Must specify model and download_dir")
 
+    if not os.path.exists(args.download_dir):
+        os.makedirs(args.download_dir)
+
     prepare_hf_model_weights(
         model_name_or_path = args.model,
         cache_dir=args.download_dir,