diff --git a/docker/build_image.py b/docker/build_image.py
index 15bfc1af..3e4fda0e 100644
--- a/docker/build_image.py
+++ b/docker/build_image.py
@@ -360,7 +360,7 @@ class SwiftImageBuilder(LLMImageBuilder):
             extra_content = extra_content.replace('{python_version}',
                                                   self.args.python_version)
         extra_content += """
-RUN pip install --no-cache-dir -U "deepspeed==0.16.*" --no-deps && \
+RUN pip install --no-cache-dir -U deepspeed --no-deps && \
     pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
 """
         version_args = (