diff --git a/CMakeLists.txt b/CMakeLists.txt
index e9c6c5c..8b1eba0 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -50,6 +50,9 @@ endif()
 set(USE_TRITONSERVER_DATATYPE "ON")
 message("-- Enable USE_TRITONSERVER_DATATYPE")
 
+# Required for Ubuntu 22.04/Triton 22.05+
+set(CMAKE_CUDA_FLAGS "-include stdio.h")
+
 set(CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake/Modules)
 
 # Python.h needed by torch headers.
diff --git a/LEGAL.md b/LEGAL.md
new file mode 100644
index 0000000..e52a5f1
--- /dev/null
+++ b/LEGAL.md
@@ -0,0 +1 @@
+This application is not part of the Snowflake Service and is governed by the terms in LICENSE, unless expressly agreed to in writing. You use this application at your own risk, and Snowflake has no obligation to support your use of this application.
diff --git a/README.md b/README.md
index b892f03..c50e458 100644
--- a/README.md
+++ b/README.md
@@ -26,7 +26,7 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 -->
 
-**NOTE: Fastertransformer backend is currently undergoing restructuring. Build instructions are only tested with Triton container versions <= `23.04`**. 
+**NOTE: Fastertransformer backend is currently undergoing restructuring so might not work with all versions of Triton.** 
 
 # FasterTransformer Backend
 
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 369ba9d..bfbfba0 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG TRITON_VERSION=22.12
+ARG TRITON_VERSION=23.05
 ARG BASE_IMAGE=nvcr.io/nvidia/tritonserver:${TRITON_VERSION}-py3
 FROM ${BASE_IMAGE}
 
-RUN apt-get update
-RUN apt-get install -y --no-install-recommends \
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
         autoconf \
         autogen \
         clangd \
@@ -35,22 +35,22 @@ RUN apt-get install -y --no-install-recommends \
         unzip \
         zstd \
         zip \
-        zsh
-RUN pip3 install torch==1.12.1+cu116 -f https://download.pytorch.org/whl/torch_stable.html && \
-    pip3 install --extra-index-url https://pypi.ngc.nvidia.com regex fire tritonclient[all] && \
-    pip3 install transformers huggingface_hub tokenizers SentencePiece sacrebleu datasets tqdm omegaconf rouge_score && \
-    pip3 install cmake==3.24.3
-
-RUN apt-get clean && \
+        zsh && \
+    apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
+RUN pip3 install --no-cache-dir --extra-index-url https://download.pytorch.org/whl/cu118 torch==2.0.1+cu118 && \
+    pip3 install --no-cache-dir --extra-index-url https://pypi.ngc.nvidia.com regex fire tritonclient[all] && \
+    pip3 install --no-cache-dir accelerate transformers huggingface_hub tokenizers SentencePiece sacrebleu datasets tqdm omegaconf rouge_score && \
+    pip3 install --no-cache-dir cmake==3.24.3
+
 # backend build
 ADD . /workspace/build/fastertransformer_backend
 RUN mkdir -p /workspace/build/fastertransformer_backend/build
 
 WORKDIR /workspace/build/fastertransformer_backend/build
 ARG FORCE_BACKEND_REBUILD=0
-RUN cmake \
+RUN CUDAFLAGS="-include stdio.h" cmake \
       -D CMAKE_EXPORT_COMPILE_COMMANDS=1 \
       -D CMAKE_BUILD_TYPE=Release \
       -D ENABLE_FP8=OFF \
@@ -58,11 +58,11 @@ RUN cmake \
       -D TRITON_COMMON_REPO_TAG="r${NVIDIA_TRITON_SERVER_VERSION}" \
       -D TRITON_CORE_REPO_TAG="r${NVIDIA_TRITON_SERVER_VERSION}" \
       -D TRITON_BACKEND_REPO_TAG="r${NVIDIA_TRITON_SERVER_VERSION}" \
-      ..
-RUN cd _deps/repo-ft-src/ && \
+      .. && \
+    cd _deps/repo-ft-src/ && \
     git log | head -n 3 2>&1 | tee /workspace/build/fastertransformer_backend/FT_version.txt && \
     cd /workspace/build/fastertransformer_backend/build && \
-    make -j"$(grep -c ^processor /proc/cpuinfo)" install && \
+    CUDAFLAGS="-include stdio.h" make -O -j"$(grep -c ^processor /proc/cpuinfo)" install && \
     rm /workspace/build/fastertransformer_backend/build/bin/*_example -rf && \
     rm /workspace/build/fastertransformer_backend/build/lib/lib*Backend.so -rf
 
diff --git a/docker/create_dockerfile_and_build.py b/docker/create_dockerfile_and_build.py
index 04309af..205f653 100644
--- a/docker/create_dockerfile_and_build.py
+++ b/docker/create_dockerfile_and_build.py
@@ -61,7 +61,7 @@ def create_dependencies(base_image):
 ARG BASE_IMAGE={base_image}
     '''.format(base_image=base_image)
     df += '''
-FROM ${BASE_IMAGE}
+FROM ${BASE_IMAGE} AS build_image
 RUN apt-get update && apt-get install -y --no-install-recommends \\
         autoconf \\
         autogen \\