feat: support finding libraries on aarch64

Signed-off-by: Naren Dasan <naren@narendasan.com>
pytorch · Jul 22, 2022 · 2abe4ba · 2abe4ba
1 parent 7680456
commit 2abe4ba
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/py/torch_tensorrt/__init__.py b/py/torch_tensorrt/__init__.py
@@ -60,7 +60,10 @@ def _find_lib(name, paths):
     elif sys.platform.startswith("linux"):
         LINUX_PATHS = [
             "/usr/local/cuda/lib64",
-        ] + os.environ["LD_LIBRARY_PATH"].split(os.path.pathsep)
+        ]
+
+        if "LD_LIBRARY_PATH" in os.environ:
+            LINUX_PATHS += os.environ["LD_LIBRARY_PATH"].split(os.path.pathsep)
 
         if platform.uname().processor == "x86_64":
             LINUX_PATHS += [