diff --git a/comps/finetuning/docker/Dockerfile_cpu b/comps/finetuning/docker/Dockerfile_cpu
index d5490c57d..afcf36f15 100644
--- a/comps/finetuning/docker/Dockerfile_cpu
+++ b/comps/finetuning/docker/Dockerfile_cpu
@@ -1,3 +1,6 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
 # Use the same python version with ray
 FROM python:3.10.14
 
@@ -5,8 +8,6 @@ ARG HF_TOKEN
 
 ENV HF_TOKEN=$HF_TOKEN
 
-RUN apt-get update -y && apt-get install -y vim htop net-tools dnsutils
-
 RUN useradd -m -s /bin/bash user && \
     mkdir -p /home/user && \
     chown -R user /home/user/
@@ -19,11 +20,11 @@ USER user
 
 ENV PATH=$PATH:/home/user/.local/bin
 
-RUN pip install --no-cache-dir --upgrade pip && \
-    python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu && \
-    python -m pip install intel-extension-for-pytorch && \
-    python -m pip install oneccl_bind_pt --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/cpu/us/ && \
-    pip install --no-cache-dir -r /home/user/comps/finetuning/requirements.txt
+RUN python -m pip install --no-cache-dir --upgrade pip && \
+    python -m pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu && \
+    python -m pip install --no-cache-dir intel-extension-for-pytorch && \
+    python -m pip install --no-cache-dir oneccl_bind_pt --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/cpu/us/ && \
+    python -m pip install --no-cache-dir -r /home/user/comps/finetuning/requirements.txt
 
 ENV PYTHONPATH=$PYTHONPATH:/home/user
 
diff --git a/comps/finetuning/docker/Dockerfile_hpu b/comps/finetuning/docker/Dockerfile_hpu
index 1277d76c1..8c1484c62 100644
--- a/comps/finetuning/docker/Dockerfile_hpu
+++ b/comps/finetuning/docker/Dockerfile_hpu
@@ -1,10 +1,12 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
 # Use the same python version with ray
-FROM vault.habana.ai/gaudi-docker/1.16.1/ubuntu22.04/habanalabs/pytorch-installer-2.2.2:latest
+# FROM vault.habana.ai/gaudi-docker/1.16.1/ubuntu22.04/habanalabs/pytorch-installer-2.2.2:latest as hpu
+FROM opea/habanalabs:1.16.1-pytorch-installer-2.2.2 as hpu
 
 ENV DEVICE="hpu"
 
-RUN apt-get update -y && apt-get install -y vim htop net-tools dnsutils
-
 RUN useradd -m -s /bin/bash user && \
     mkdir -p /home/user && \
     chown -R user /home/user/
@@ -17,15 +19,12 @@ USER user
 
 ENV PATH=$PATH:/home/user/.local/bin
 
-RUN pip install --no-cache-dir --upgrade pip && \
-    pip install --no-cache-dir -r /home/user/comps/finetuning/requirements.txt && \
-    pip install --no-cache-dir optimum-habana
+RUN python -m pip install --no-cache-dir --upgrade pip && \
+    python -m pip install --no-cache-dir -r /home/user/comps/finetuning/requirements.txt && \
+    python -m pip install --no-cache-dir optimum-habana
 
 ENV PYTHONPATH=$PYTHONPATH:/home/user
 
 WORKDIR /home/user/comps/finetuning
 
 ENTRYPOINT ["/bin/bash", "launch.sh"]
-
-# CMD ["/bin/bash"]
-