diff --git a/Dockerfile b/Dockerfile
index 9c88dd4..8b21393 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@fix_scaling#egg=runpod --compile
+RUN pip install git+https://github.com/runpod/runpod-python@main#egg=runpod --compile
 
 # Prepare the models inside the docker image
 ARG HUGGING_FACE_HUB_TOKEN=
diff --git a/builder/requirements.txt b/builder/requirements.txt
index 9b2aee3..a124420 100644
--- a/builder/requirements.txt
+++ b/builder/requirements.txt
@@ -3,4 +3,4 @@
 fastapi==0.99.1
 vllm==0.1.2
 huggingface-hub==0.16.4
-runpod @ git+https://github.com/runpod/runpod-python@fix_scaling#egg=runpod 
+runpod @ git+https://github.com/runpod/runpod-python@main#egg=runpod