diff --git a/Dockerfile b/Dockerfile
index 3292e25..28f2e58 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -6,7 +6,7 @@ RUN dnf -y install --setopt=install_weak_deps=False --nodocs \
   gcc gcc-c++ binutils make ffmpeg-free \
   cmake ninja-build aria2c tar xz vim nano \
   libdrm-devel zlib-devel openssl-devel jq \
-  numactl-devel gperftools-libs dialog procps-ng \
+  numactl-devel gperftools-libs procps-ng \
   && dnf clean all && rm -rf /var/cache/dnf/*
 
 # 2. Install "TheRock" ROCm SDK (Tarball Method)
@@ -54,16 +54,17 @@ ENV PIP_NO_CACHE_DIR=1
 RUN printf 'source /opt/venv/bin/activate\n' > /etc/profile.d/venv.sh
 RUN python -m pip install --upgrade pip wheel packaging "setuptools<80.0.0"
 
-# 5. Install PyTorch (TheRock Nightly)
+# 5. Install PyTorch (TheRock Nightly) and PyYAML
 RUN python -m pip install \
   --index-url https://rocm.nightlies.amd.com/v2-staging/gfx120X-all/ \
-  --pre torch torchaudio torchvision
+  --pre torch torchaudio torchvision && \
+  python -m pip install pyyaml
 
 # Flash-Attention
 WORKDIR /opt
 ENV FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"
 
-RUN git clone https://github.com/ROCm/flash-attention.git &&\ 
+RUN git clone https://github.com/ROCm/flash-attention.git &&\
   cd flash-attention &&\
   git checkout main_perf &&\
   python setup.py install && \
@@ -117,7 +118,8 @@ ENV CXX="/opt/rocm/llvm/bin/clang++"
 
 RUN export HIP_DEVICE_LIB_PATH=$(find /opt/rocm -type d -name bitcode -print -quit) && \
   echo "Compiling with Bitcode: $HIP_DEVICE_LIB_PATH" && \
-  export CMAKE_ARGS="-DROCM_PATH=/opt/rocm -DHIP_PATH=/opt/rocm -DAMDGPU_TARGETS=gfx1201 -DHIP_ARCHITECTURES=gfx1201" && \
+  export CMAKE_PREFIX_PATH="/opt/venv/lib64/python3.13/site-packages/torch/share/cmake:/opt/rocm" && \
+  export CMAKE_ARGS="-DROCM_PATH=/opt/rocm -DHIP_PATH=/opt/rocm -DAMDGPU_TARGETS=gfx1201 -DHIP_ARCHITECTURES=gfx1201 -DCMAKE_PREFIX_PATH=/opt/venv/lib64/python3.13/site-packages/torch/share/cmake:/opt/rocm" && \
   python -m pip wheel --no-build-isolation --no-deps -w /tmp/dist -v . && \
   python -m pip install /tmp/dist/*.whl
 
diff --git a/build_and_run.sh b/build_and_run.sh
index db5b18f..c07d3d4 100644
--- a/build_and_run.sh
+++ b/build_and_run.sh
@@ -261,10 +261,11 @@ build_image() {
     print_step "开始构建 Docker 镜像..."
     print_info "镜像名称：${IMAGE_NAME}:${IMAGE_TAG}"
     print_info "构建目录：${PROJECT_DIR}"
+    print_info "使用缓存策略：未修改的层将使用缓存"
     
     BUILD_START=$(date +%s)
     
-    if docker build --no-cache -t "${IMAGE_NAME}:${IMAGE_TAG}" .; then
+    if docker build -t "${IMAGE_NAME}:${IMAGE_TAG}" .; then
         BUILD_END=$(date +%s)
         BUILD_TIME=$((BUILD_END - BUILD_START))