Fix jetson onnxruntime (#14698)

* Fix jetson onnxruntime * Remove comment
blakeblackshear · Oct 31, 2024 · bb4e863 · bb4e863
1 parent c7a4220
commit bb4e863
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/docker/tensorrt/Dockerfile.arm64 b/docker/tensorrt/Dockerfile.arm64
@@ -41,8 +41,11 @@ RUN --mount=type=bind,source=docker/tensorrt/detector/build_python_tensorrt.sh,t
     && TENSORRT_VER=$(cat /etc/TENSORRT_VER) /deps/build_python_tensorrt.sh
 
 COPY docker/tensorrt/requirements-arm64.txt /requirements-tensorrt.txt
-RUN pip3 uninstall -y onnxruntime \
-    && pip3 wheel --wheel-dir=/trt-wheels -r /requirements-tensorrt.txt
+ADD https://nvidia.box.com/shared/static/9aemm4grzbbkfaesg5l7fplgjtmswhj8.whl /tmp/onnxruntime_gpu-1.15.1-cp39-cp39-linux_aarch64.whl
+
+RUN pip3 uninstall -y onnxruntime-openvino \
+    && pip3 wheel --wheel-dir=/trt-wheels -r /requirements-tensorrt.txt \
+    && pip3 install --no-deps /tmp/onnxruntime_gpu-1.15.1-cp39-cp39-linux_aarch64.whl
 
 FROM build-wheels AS trt-model-wheels
 ARG DEBIAN_FRONTEND

diff --git a/docker/tensorrt/requirements-arm64.txt b/docker/tensorrt/requirements-arm64.txt
@@ -1,2 +1 @@
-cuda-python == 11.7; platform_machine == 'aarch64'
-onnxruntime @ https://nvidia.box.com/shared/static/9aemm4grzbbkfaesg5l7fplgjtmswhj8.whl; platform_machine == 'aarch64'
+cuda-python == 11.7; platform_machine == 'aarch64'