diff --git a/src/engine.py b/src/engine.py
index 77544b6..9cfed73 100644
--- a/src/engine.py
+++ b/src/engine.py
@@ -21,7 +21,7 @@ class vLLMEngine:
     def __init__(self, engine = None):
         load_dotenv() # For local development
         self.config = EngineConfig().config
-        self.tokenizer = TokenizerWrapper(self.config.get("tokenizer"), self.config.get("tokenizer_revision"))
+        self.tokenizer = TokenizerWrapper(self.config.get("tokenizer"), self.config.get("tokenizer_revision"), self.config.get("trust_remote_code"))
         self.llm = self._initialize_llm() if engine is None else engine
         self.max_concurrency = int(os.getenv("MAX_CONCURRENCY", DEFAULT_MAX_CONCURRENCY))
         self.default_batch_size = int(os.getenv("DEFAULT_BATCH_SIZE", DEFAULT_BATCH_SIZE))
diff --git a/src/tokenizer.py b/src/tokenizer.py
index dfe848a..62a4cbf 100644
--- a/src/tokenizer.py
+++ b/src/tokenizer.py
@@ -3,8 +3,8 @@ import os
 from typing import Union
 
 class TokenizerWrapper:
-    def __init__(self, tokenizer_name_or_path, tokenizer_revision):
-        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, revision=tokenizer_revision)
+    def __init__(self, tokenizer_name_or_path, tokenizer_revision, trust_remote_code):
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, revision=tokenizer_revision, trust_remote_code=trust_remote_code)
         self.custom_chat_template = os.getenv("CUSTOM_CHAT_TEMPLATE")
         self.has_chat_template = bool(self.tokenizer.chat_template) or bool(self.custom_chat_template)
         if self.custom_chat_template and isinstance(self.custom_chat_template, str):
diff --git a/vllm-base/Dockerfile b/vllm-base/Dockerfile
index 05d126f..2ea7224 100644
--- a/vllm-base/Dockerfile
+++ b/vllm-base/Dockerfile
@@ -6,7 +6,7 @@
 ##########################################################
 
 # Define the CUDA version for the build
-ARG WORKER_CUDA_VERSION=12.1.0
+ARG WORKER_CUDA_VERSION=11.8.0
 
 FROM nvidia/cuda:${WORKER_CUDA_VERSION}-devel-ubuntu22.04 AS dev
 
@@ -76,7 +76,7 @@ ENV NVCC_THREADS=${nvcc_threads}
 # Build extensions
 RUN python3 setup.py build_ext --inplace
 
-FROM nvidia/cuda:${WORKER_CUDA_VERSION}-base-ubuntu22.04 AS vllm-base
+FROM nvidia/cuda:${WORKER_CUDA_VERSION}-runtime-ubuntu22.04 AS vllm-base
 
 # Re-declare ARG after FROM
 ARG WORKER_CUDA_VERSION