diff --git a/dev/docker/Dockerfile.habana b/dev/docker/Dockerfile.habana
new file mode 100644
index 000000000..508361d22
--- /dev/null
+++ b/dev/docker/Dockerfile.habana
@@ -0,0 +1,10 @@
+FROM vault.habana.ai/gaudi-docker/1.14.0/ubuntu22.04/habanalabs/pytorch-installer-2.1.1:latest
+
+ENV LANG=en_US.UTF-8
+
+COPY ../../ /root/llm-on-ray
+
+WORKDIR /root/llm-on-ray
+
+RUN pip install . && \
+    pip install --upgrade-strategy eager optimum[habana]
\ No newline at end of file
diff --git a/docs/setup.md b/docs/setup.md
index 5da340628..5748d9785 100644
--- a/docs/setup.md
+++ b/docs/setup.md
@@ -29,6 +29,8 @@ Intel® 1st, 2nd, 3rd, and 4th Gen Xeon® Scalable Performance processor
 #### 1. Prerequisites
 For Intel GPU, ensure the [Intel® oneAPI Base Toolkit](https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit.html) is installed.
 
+For Gaudi, ensure the [SynapseAI SW stack and container runtime](https://docs.habana.ai/en/latest/Installation_Guide/Bare_Metal_Fresh_OS.html?highlight=installer#run-using-containers) is installed.
+
 #### 2. Clone the repository and install dependencies.
 ```bash
 git clone https://github.com/intel/llm-on-ray.git
@@ -51,7 +53,7 @@ source $(python -c "import oneccl_bindings_for_pytorch as torch_ccl;print(torch_
 
 For Gaudi:
 
-Please use the [Dockerfile](../inference/habana/Dockerfile) to build the image. Alternatively, you can install the dependecies on a bare metal machine. In this case, please refer to [here](https://docs.habana.ai/en/latest/Installation_Guide/Bare_Metal_Fresh_OS.html#build-docker-bare).
+Please use the [Dockerfile](../dev/docker/Dockerfile.habana) to build the image. Alternatively, you can install the dependecies on a bare metal machine. In this case, please refer to [here](https://docs.habana.ai/en/latest/Installation_Guide/Bare_Metal_Fresh_OS.html#build-docker-bare).
 
 ```bash
 docker build \
diff --git a/inference/habana/Dockerfile b/inference/habana/Dockerfile
deleted file mode 100644
index c92430b80..000000000
--- a/inference/habana/Dockerfile
+++ /dev/null
@@ -1,16 +0,0 @@
-FROM vault.habana.ai/gaudi-docker/1.12.0/ubuntu22.04/habanalabs/pytorch-installer-2.0.1:latest
-
-ENV LANG=en_US.UTF-8
-ENV PYTHONPATH=/root:/usr/lib/habanalabs/
-
-RUN cd ~ && \
-    git clone https://github.com/huggingface/optimum-habana.git && \
-    cd optimum-habana/ && git reset --hard b6edce65b70e0fadd5d5f51234700bd1144cd0b0 && pip install -e . && cd ../ && \
-    cd ./optimum-habana/examples/text-generation/ && \
-    pip install -r requirements.txt && \
-    cd ~
-
-RUN pip install https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp310-cp310-manylinux2014_x86_64.whl && \
-    pip install ray[tune,serve] transformers==4.32.0 accelerate==0.22.0 pydantic_yaml
-
-WORKDIR /root/llm-ray
\ No newline at end of file