add library_dirs (#977)

2023-06-15 00:09:12 +08:00 · 2023-06-15 00:09:12 +08:00 · f6d42f2dd0
commit f6d42f2dd0
parent 473a67073e
1 changed files with 5 additions and 0 deletions
--- a/python/setup.py
+++ b/python/setup.py
@ -108,10 +108,15 @@ include_dirs = [
 ]


+library_dirs = [
+    cuda_install_path + '/lib64',
+]
+
 ext_modules = [
    Pybind11Extension('cutlass_bindings',
                      ['cutlass/cpp/cutlass_bindings.cpp'],
                      include_dirs=include_dirs,
+                      library_dirs=library_dirs,
                      extra_compile_args=['-Xcompiler="-fpermissive"', '-w', '-std=c++17'],
                      libraries=['cudart'])
 ]