diff --git a/comfyui-rocm-dockerfile b/comfyui-rocm-dockerfile
index 1f6d5aa..5e33313 100644
--- a/comfyui-rocm-dockerfile
+++ b/comfyui-rocm-dockerfile
@@ -7,15 +7,9 @@ ENV PYTHONUNBUFFERED=1
 ENV DEBIAN_FRONTEND=noninteractive
 
 # Set path variables
-ENV Torch_DIR=/usr/lib/python3/dist-packages/torch/share/cmake/Torch
-ENV CMAKE_PREFIX_PATH=/usr/lib/python3/dist-packages/torch/share/cmake
 ENV ROCM_PATH=/usr
 ENV HIP_PATH=/usr
 
-# Set environment variables
-ENV PYTORCH_ROCM_ARCH=gfx1100
-ENV HIP_VISIBLE_DEVICES=0
-
 # Install libraries
 RUN apt update -y && apt install -y --no-install-recommends \
     rocm \
@@ -48,6 +42,15 @@ WORKDIR /builder
 RUN git clone --depth 1 --branch v2.9.1 https://github.com/pytorch/audio.git && \
     cd audio && \
     git submodule update --init --recursive && \
+    # Set path variables
+    export CMAKE_PREFIX_PATH=/usr/lib/python3/dist-packages/torch/share/cmake && \
+    export Torch_DIR=/usr/lib/python3/dist-packages/torch/share/cmake/Torch && \
+    export hip_DIR=/usr/lib/aarch64-linux-gnu/cmake/hip && \
+    export CMAKE_PREFIX_PATH="/usr/lib/python3/dist-packages/torch/share/cmake:/usr/lib/aarch64-linux-gnu/cmake/hip" && \
+    # Set environment variables
+    export PYTORCH_ROCM_ARCH=gfx1100 && \
+    export HIP_VISIBLE_DEVICES=0 && \
+    # Execute build
     USE_ROCM=1 BUILD_ROCM=1 CC=gcc CXX=g++ python3 setup.py install && \
     cd .. && rm -rf audio