From 12d855003ba92c8a15d1739ce65a14c6fb16e254 Mon Sep 17 00:00:00 2001 From: Shubham <32878682+IamShubhamGupto@users.noreply.github.com> Date: Fri, 22 Mar 2024 23:47:47 -0400 Subject: [PATCH] Add support for Nvidia Jetson hw (#262) https://developer.nvidia.com/cuda-gpus --- setup.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index d794a9ce..c17ab0bb 100644 --- a/setup.py +++ b/setup.py @@ -104,11 +104,19 @@ def append_nvcc_threads(nvcc_extra_args): f"{PACKAGE_NAME} is only supported on CUDA 11.6 and above. " "Note: make sure nvcc has a supported version by running nvcc -V." ) - + + cc_flag.append("-gencode") + cc_flag.append("arch=compute_53,code=sm_53") + cc_flag.append("-gencode") + cc_flag.append("arch=compute_62,code=sm_62") cc_flag.append("-gencode") cc_flag.append("arch=compute_70,code=sm_70") cc_flag.append("-gencode") + cc_flag.append("arch=compute_72,code=sm_72") + cc_flag.append("-gencode") cc_flag.append("arch=compute_80,code=sm_80") + cc_flag.append("-gencode") + cc_flag.append("arch=compute_87,code=sm_87") if bare_metal_version >= Version("11.8"): cc_flag.append("-gencode") cc_flag.append("arch=compute_90,code=sm_90")