diff --git a/Dockerfile_amd b/Dockerfile_amd
index 68561bdf..6f8f874b 100644
--- a/Dockerfile_amd
+++ b/Dockerfile_amd
@@ -36,7 +36,7 @@ COPY launcher launcher
 RUN cargo build --release
 
 # Text Generation Inference base image for RoCm
-FROM rocm/dev-ubuntu-22.04:6.1 as base
+FROM rocm/dev-ubuntu-22.04:6.1.1_hip_update as base
 
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
     build-essential \
@@ -114,13 +114,6 @@ ARG BUILD_CAFFE2="0" \
 
 RUN cd pytorch && python tools/amd_build/build_amd.py && python setup.py install
 
-ARG GITHUB_TOKEN
-RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends wget && \
-    rm -rf /var/lib/apt/lists/* && \
-    wget --header "Authorization: token ${GITHUB_TOKEN}" https://raw.githubusercontent.com/fxmarty/patched_hipruntime/main/libamdhip64.so.6
-
-ENV LD_PRELOAD="/libamdhip64.so.6"
-
 # Set as recommended: https://github.com/ROCm/triton/wiki/A-script-to-set-program-execution-environment-in-ROCm
 ENV HIP_FORCE_DEV_KERNARG=1
 
@@ -193,6 +186,10 @@ COPY --from=exllamav2-kernels-builder /usr/src/build/lib.linux-x86_64-cpython-31
 COPY proto proto
 COPY server server
 COPY server/Makefile server/Makefile
+RUN cd server && \
+    make gen-server && \
+    pip install -r requirements_rocm.txt && \
+    pip install ".[accelerate, peft, outlines]" --no-cache-dir
 
 # Install benchmarker
 COPY --from=builder /usr/src/target/release/text-generation-benchmark /usr/local/bin/text-generation-benchmark
@@ -201,11 +198,6 @@ COPY --from=builder /usr/src/target/release/text-generation-router /usr/local/bi
 # Install launcher
 COPY --from=builder /usr/src/target/release/text-generation-launcher /usr/local/bin/text-generation-launcher
 
-RUN cd server && \
-    make gen-server && \
-    pip install -r requirements_rocm.txt
-    #pip install ".[accelerate, peft, outlines]" --no-cache-dir
-
 # AWS Sagemaker compatible image
 FROM base as sagemaker
 
@@ -217,8 +209,8 @@ ENTRYPOINT ["./entrypoint.sh"]
 # Final image
 FROM base-copy
 
-# COPY ./tgi-entrypoint.sh /tgi-entrypoint.sh
-# RUN chmod +x /tgi-entrypoint.sh
+COPY ./tgi-entrypoint.sh /tgi-entrypoint.sh
+RUN chmod +x /tgi-entrypoint.sh
 
-# ENTRYPOINT ["/tgi-entrypoint.sh"]
-# CMD ["--json-output"]
+ENTRYPOINT ["/tgi-entrypoint.sh"]
+CMD ["--json-output"]