diff --git a/Makefile b/Makefile
index 1ae0aa514..02e82be73 100644
--- a/Makefile
+++ b/Makefile
@@ -40,19 +40,15 @@ PACKAGE_FILES = $(PACKAGE_PYTHON_FILES)  \
 $(PACKAGE_DIST) $(PACKAGE_WHEEL): $(PACKAGE_FILES)
 	python -m build
 
-TGI_VERSION ?= 2.1.1
-
 neuronx-tgi: $(PACKAGE_DIST)
 	docker build --rm -f text-generation-inference/Dockerfile \
 	             --build-arg VERSION=$(VERSION) \
-	             --build-arg TGI_VERSION=$(TGI_VERSION) \
 				 -t neuronx-tgi:$(VERSION) .
 	docker tag neuronx-tgi:$(VERSION) neuronx-tgi:latest
 
 neuronx-tgi-sagemaker: $(PACKAGE_DIST)
 	docker build --rm -f text-generation-inference/Dockerfile \
 	             --build-arg VERSION=$(VERSION) \
-	             --build-arg TGI_VERSION=$(TGI_VERSION) \
 				 --target sagemaker \
 				 -t neuronx-tgi:$(VERSION) .
 
@@ -90,7 +86,7 @@ test_installs:
 tgi_server:
 	python -m pip install -r text-generation-inference/server/build-requirements.txt
 	make -C text-generation-inference/server clean
-	VERSION=${VERSION} TGI_VERSION=${TGI_VERSION} make -C text-generation-inference/server gen-server
+	VERSION=${VERSION} make -C text-generation-inference/server gen-server
 
 tgi_test: tgi_server
 	python -m pip install .[neuronx]
diff --git a/text-generation-inference/Dockerfile b/text-generation-inference/Dockerfile
index ac23d637e..8dab8b378 100644
--- a/text-generation-inference/Dockerfile
+++ b/text-generation-inference/Dockerfile
@@ -1,7 +1,6 @@
-# Fetch and extract the TGI sources (TGI_VERSION is mandatory)
+# Fetch and extract the TGI sources
 FROM alpine AS tgi
-ARG TGI_VERSION
-RUN test -n ${TGI_VERSION:?}
+ARG TGI_VERSION=2.1.1
 RUN mkdir -p /tgi
 ADD https://github.com/huggingface/text-generation-inference/archive/refs/tags/v${TGI_VERSION}.tar.gz /tgi/sources.tar.gz
 RUN tar -C /tgi -xf /tgi/sources.tar.gz --strip-components=1
@@ -13,7 +12,7 @@ WORKDIR /usr/src
 
 ARG CARGO_REGISTRIES_CRATES_IO_PROTOCOL=sparse
 
-FROM chef as planner
+FROM chef AS planner
 COPY --from=tgi /tgi/Cargo.lock Cargo.lock
 COPY --from=tgi /tgi/Cargo.toml Cargo.toml
 COPY --from=tgi /tgi/rust-toolchain.toml rust-toolchain.toml
@@ -100,6 +99,7 @@ RUN apt-get update -y \
     aws-neuronx-collectives=2.22.26.0-17a033bc8 \
     aws-neuronx-runtime-lib=2.22.14.0-6e27b8d5b \
     aws-neuronx-tools=2.19.0.0 \
+    libxml2 \
     && rm -rf /var/lib/apt/lists/* \
     && apt-get clean
 
@@ -117,7 +117,7 @@ RUN pip3 install \
     hf_transfer huggingface_hub
 
 # Install optimum-neuron
-COPY dist/optimum_neuron-${VERSION}.tar.gz optimum-neuron.tar.gz
+COPY ./dist/optimum_neuron-${VERSION}.tar.gz optimum-neuron.tar.gz
 RUN pip3 install optimum-neuron.tar.gz
 
 # TGI base env
@@ -134,7 +134,7 @@ COPY --from=pyserver /pyserver/build/dist dist
 RUN pip install dist/text_generation_server*.tar.gz
 
 # AWS Sagemaker compatible image
-FROM neuron as sagemaker
+FROM neuron AS sagemaker
 
 COPY text-generation-inference/sagemaker-entrypoint.sh entrypoint.sh
 RUN chmod +x entrypoint.sh
diff --git a/text-generation-inference/server/Makefile b/text-generation-inference/server/Makefile
index 2637f01d1..9b748e48f 100644
--- a/text-generation-inference/server/Makefile
+++ b/text-generation-inference/server/Makefile
@@ -2,7 +2,7 @@
 pkg_name := text_generation_server
 BUILDDIR ?= $(CURDIR)/build
 VERSION ?= 0.0.1
-TGI_VERSION ?= 2.0.2
+TGI_VERSION ?= 2.1.1
 mkfile_path := $(abspath $(lastword $(MAKEFILE_LIST)))
 mkfile_dir := $(dir $(mkfile_path))
 pkg_dir := $(BUILDDIR)/$(pkg_name)