[aarch64] set rpath for libs copied for cuda sbsa wheel (#1920)

tinglvv · web-flow · commit 1914e7678bb3 · 2024-07-19T15:13:00.000-04:00
* set rpath for libs copied for arm

* fix lint
diff --git a/aarch64_linux/aarch64_wheel_ci_build.py b/aarch64_linux/aarch64_wheel_ci_build.py
@@ -95,12 +95,10 @@ def update_wheel(wheel_path) -> None:
     for lib_path in libs_to_copy:
         lib_name = os.path.basename(lib_path)
         shutil.copy2(lib_path, f"{folder}/tmp/torch/lib/{lib_name}")
-    os.system(
-        f"cd {folder}/tmp/torch/lib/; patchelf --set-rpath '$ORIGIN' {folder}/tmp/torch/lib/libtorch_cuda.so"
-    )
-    os.system(
-        f"cd {folder}/tmp/torch/lib/; patchelf --set-rpath '$ORIGIN' {folder}/tmp/torch/lib/libcudnn_graph.so.9"
-    )
+        os.system(
+            f"cd {folder}/tmp/torch/lib/; "
+            f"patchelf --set-rpath '$ORIGIN' --force-rpath {folder}/tmp/torch/lib/{lib_name}"
+        )
     os.mkdir(f"{folder}/cuda_wheel")
     os.system(f"cd {folder}/tmp/; zip -r {folder}/cuda_wheel/{wheelname} *")
     shutil.move(
