diff --git a/Dockerfile_amd b/Dockerfile_amd
index 8bbcbbe8..f1c68674 100644
--- a/Dockerfile_amd
+++ b/Dockerfile_amd
@@ -96,7 +96,6 @@ RUN pip uninstall -y triton && \
     pip install .
 
 RUN git clone --depth 1 --recursive --single-branch --branch 2.3-patched https://github.com/fxmarty/pytorch.git pytorch && cd pytorch && pip install -r requirements.txt --no-cache-dir
-# RUN git clone --depth 1 --recursive --single-branch --branch release/2.3 https://github.com/pytorch/pytorch.git pytorch && cd pytorch && pip install -r requirements.txt --no-cache-dir
 
 ARG _GLIBCXX_USE_CXX11_ABI="1"
 ARG CMAKE_PREFIX_PATH="/opt/conda"
@@ -116,6 +115,7 @@ ARG BUILD_CAFFE2="0" \
 RUN cd pytorch && python tools/amd_build/build_amd.py && python setup.py install
 
 # Set as recommended: https://github.com/ROCm/triton/wiki/A-script-to-set-program-execution-environment-in-ROCm
+# Disabled for now as it is currently not stable with ROCm 6.1.
 # ENV HIP_FORCE_DEV_KERNARG=1
 
 FROM base AS kernel-builder
@@ -184,8 +184,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
-    # pip install -r requirements_rocm.txt && \
-    #pip install ".[accelerate, peft, outlines]" --no-cache-dir
+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
@@ -194,21 +196,19 @@ 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
-
 # AWS Sagemaker compatible image
-FROM base-copy as sagemaker
+FROM base as sagemaker
+
 COPY sagemaker-entrypoint.sh entrypoint.sh
 RUN chmod +x entrypoint.sh
 
 ENTRYPOINT ["./entrypoint.sh"]
 
 # Final image
-FROM base-copy
+FROM base
 
-# ENTRYPOINT ["text-generation-launcher"]
-# CMD ["--json-output"]
+COPY ./tgi-entrypoint.sh /tgi-entrypoint.sh
+RUN chmod +x /tgi-entrypoint.sh
 
-# NOTE: Temporarily, for TGI, please mount a volume and install locally the server with `cd /tgi/server && pip install ".[accelerate, peft, outlines]" --no-cache-dir`
\ No newline at end of file
+ENTRYPOINT ["/tgi-entrypoint.sh"]
+CMD ["--json-output"]