diff --git a/Dockerfile b/Dockerfile
index eabbfd4af..9de139a0f 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -3,6 +3,8 @@
 
 FROM nvidia/cuda:13.3.0-devel-ubuntu24.04@sha256:ef2203909e80b8b976cfc672f7e2ae2b00bc0e25c404ee86d89e10a3802f1c52
 
+ARG ARCH=aarch64
+
 ENV DEBIAN_FRONTEND=noninteractive
 ENV TORCH_CUDA_ARCH_LIST=10.0
 ENV OMPI_ALLOW_RUN_AS_ROOT=1
@@ -81,7 +83,7 @@ RUN python3.12 -m venv "$VIRTUAL_ENV" \
 ENV TRT_LIB_DIR=/opt/venv/lib/python3.12/site-packages/tensorrt_libs
 ENV NCCL_LIB_DIR=/opt/venv/lib/python3.12/site-packages/nvidia/nccl/lib
 ENV TVM_FFI_LIB_DIR=/opt/venv/lib/python3.12/site-packages/tvm_ffi/lib
-ENV TRT_INC_DIR=/usr/include/aarch64-linux-gnu
+ENV TRT_INC_DIR=/usr/include/$ARCH-linux-gnu
 ENV LD_LIBRARY_PATH=$TRT_LIB_DIR:$NCCL_LIB_DIR:$TVM_FFI_LIB_DIR:/usr/local/cuda/lib64
 ENV LD_PRELOAD=/usr/local/cuda/lib64/libcublas.so.13
 
@@ -92,7 +94,7 @@ RUN python3.12 -c \
     && test -f "$TRT_LIB_DIR/libnvonnxparser.so.11" \
     && test -f "$NCCL_LIB_DIR/libnccl.so.2" \
     && test -f "$TVM_FFI_LIB_DIR/libtvm_ffi.so" \
-    && mpirun --tag-output -np 1 true
+    && HWLOC_COMPONENTS=-gl mpirun --tag-output -np 1 true
 
 WORKDIR /workspace/tensorrt-model-connect
 CMD ["bash"]
diff --git a/docker-compose.yml b/docker-compose.yml
index 5dcf01e6d..f7ccc8538 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -6,6 +6,8 @@ services:
     build:
       context: .
       dockerfile: Dockerfile
+      args:
+        - ARCH=${ARCH:-aarch64}
     container_name: trtmc-dev
     working_dir: /workspace/tensorrt-model-connect
     tty: true