diff --git a/TensorRT_whisper.md b/TensorRT_whisper.md
index aeed8db..37b1756 100644
--- a/TensorRT_whisper.md
+++ b/TensorRT_whisper.md
@@ -5,24 +5,33 @@ We have only tested the TensorRT backend in docker so, we recommend docker for a
 ## Installation
 - Install [docker](https://docs.docker.com/engine/install/)
 - Install [nvidia-container-toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
+
 - Clone this repo.
 ```bash
 git clone https://github.com/collabora/WhisperLive.git
 cd WhisperLive
 ```
+
 - Build the TensorRT-LLM docker image
 ```bash
 docker build --file docker/Dockerfile.tensorrt --tag tensorrt_llm:latest .
 ```
 **NOTE**: This could take some time.
+
 - Next, we run the docker image and mount WhisperLive repo to the containers `/home` directory.
 ```bash
 docker run -it --gpus all --shm-size=8g \
        --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \
        -v /path/to/WhisperLive:/home/WhisperLive \
        tensorrt_llm:latest
+
+cd /home/TensorRT-LLM
+python3 scripts/build_wheel.py --cuda_architectures "89-real" --clean --trt_root /usr/local/tensorrt
+pip install build/tensorrt_llm*.whl
 ```
-- Once inside the docker container, make sure to test the installation. 
+**NOTE**: `--cuda_architectures "89-real"` builds for 4090, change according to your device.
+
+- Make sure to test the installation. 
 ```bash
 # export ENV=${ENV:-/etc/shinit_v2} 
 # source $ENV
diff --git a/docker/Dockerfile.tensorrt b/docker/Dockerfile.tensorrt
index 25fd666..00de173 100644
--- a/docker/Dockerfile.tensorrt
+++ b/docker/Dockerfile.tensorrt
@@ -1,12 +1,11 @@
 ARG BASE_IMAGE=nvcr.io/nvidia/pytorch
 ARG BASE_TAG=23.10-py3
-ARG CUDA_ARCHITECTURES="89-real"
 
 FROM ${BASE_IMAGE}:${BASE_TAG} as base
 
 WORKDIR /home
 COPY scripts/install_tensorrt_llm.sh install_tensorrt_llm.sh
-RUN bash install_tensorrt_llm.sh ${CUDA_ARCHITECTURES} && rm install_tensorrt_llm.sh
+RUN bash install_tensorrt_llm.sh && rm install_tensorrt_llm.sh
 
 COPY scripts/setup.sh /home
 RUN apt update && bash setup.sh && rm setup.sh
diff --git a/scripts/install_tensorrt_llm.sh b/scripts/install_tensorrt_llm.sh
index b74ac78..0d721db 100644
--- a/scripts/install_tensorrt_llm.sh
+++ b/scripts/install_tensorrt_llm.sh
@@ -1,9 +1,8 @@
 #!/bin/bash
 
 apt-get update && apt-get -y install git git-lfs
-git clone https://github.com/makaveli10/TensorRT-LLM.git
+git clone -b cuda12.2 https://github.com/makaveli10/TensorRT-LLM.git
 cd TensorRT-LLM
-git checkout main
 git submodule update --init --recursive
 git lfs install
 git lfs pull
@@ -16,27 +15,5 @@ bash install_cmake.sh
 bash install_ccache.sh
 bash install_tensorrt.sh
 bash install_polygraphy.sh
-source $ENV
-
-cd /root
-wget https://download.open-mpi.org/release/open-mpi/v3.1/openmpi-3.1.5.tar.gz
-tar -xzvf openmpi-3.1.5.tar.gz
-rm -rf openmpi-3.1.5.tar.gz
-cd openmpi-3.1.5
-./configure --prefix=/usr/local/
-sudo make all install
-echo 'export PATH=$PATH:/usr/local/bin' >> "${ENV}"
-echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib' >> "${ENV}"
-source $ENV
-
-cd /home/TensorRT-LLM/docker/common/
 bash install_mpi4py.sh
 source $ENV
-
-cuda_architectures="${1:-89-real}"
-
-
-cd /home/TensorRT-LLM
-python3 scripts/build_wheel.py --cuda_architectures "${cuda_architectures}" --clean --trt_root /usr/local/tensorrt
-pip install build/tensorrt_llm-*.whl
-pip install --no-deps openai-whisper
\ No newline at end of file