From db4152c4cab10decbc81eeaf3111f99a76e38e14 Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Thu, 27 Feb 2025 16:24:03 -0700
Subject: [PATCH] Fix jetson (#16854)

* Fix jetson build

* Update ci.yml

* Update Dockerfile.base

* Update Dockerfile.base

* Update Dockerfile.base

* Fix

* Update ci.yml
---
 docker/tensorrt/Dockerfile.base | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/docker/tensorrt/Dockerfile.base b/docker/tensorrt/Dockerfile.base
index 6d8d9591b..5ae018773 100644
--- a/docker/tensorrt/Dockerfile.base
+++ b/docker/tensorrt/Dockerfile.base
@@ -16,8 +16,16 @@ RUN apt-get update \
 RUN --mount=type=bind,source=docker/tensorrt/detector/tensorrt_libyolo.sh,target=/tensorrt_libyolo.sh \
     /tensorrt_libyolo.sh
 
+# COPY required individual CUDA deps
+RUN mkdir -p /usr/local/cuda-deps
+RUN if [ "$TARGETARCH" = "amd64" ]; then \
+      cp /usr/local/cuda-12.3/targets/x86_64-linux/lib/libcurand.so.* /usr/local/cuda-deps/ && \
+      cp /usr/local/cuda-12.3/targets/x86_64-linux/lib/libnvrtc.so.* /usr/local/cuda-deps/ ; \
+    fi
+
 # Frigate w/ TensorRT Support as separate image
 FROM deps AS tensorrt-base
+ARG TARGETARCH
 
 #Disable S6 Global timeout
 ENV S6_CMD_WAIT_FOR_SERVICES_MAXTIME=0
@@ -26,9 +34,8 @@ ENV S6_CMD_WAIT_FOR_SERVICES_MAXTIME=0
 COPY --from=trt-deps /usr/local/lib/libyolo_layer.so /usr/local/lib/libyolo_layer.so
 COPY --from=trt-deps /usr/local/src/tensorrt_demos /usr/local/src/tensorrt_demos
 
-# COPY Individual CUDA deps
-COPY --from=trt-deps /usr/local/cuda-12.3/targets/x86_64-linux/lib/libcurand.so.* /usr/local/cuda/
-COPY --from=trt-deps /usr/local/cuda-12.3/targets/x86_64-linux/lib/libnvrtc.so.* /usr/local/cuda/
+# COPY Individual CUDA deps folder
+COPY --from=trt-deps /usr/local/cuda-deps /usr/local/cuda
 
 COPY docker/tensorrt/detector/rootfs/ /
 ENV YOLO_MODELS=""