Use CUDA_HOME when compiling nn module

Ilyushin · Oct 30, 2020 · 052fd45 · 052fd45
1 parent 8e1afdf
commit 052fd45
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lib/csrc/nn/setup.py b/lib/csrc/nn/setup.py
@@ -1,6 +1,6 @@
 import os
 
-cuda_include='/usr/local/cuda-9.0/include'
+cuda_include=os.path.join(os.environ.get('CUDA_HOME'), 'include')
 os.system('nvcc src/nearest_neighborhood.cu -c -o src/nearest_neighborhood.cu.o -x cu -Xcompiler -fPIC -O2 -arch=sm_52 -I {}'.format(cuda_include))
 
 from cffi import FFI
@@ -16,7 +16,7 @@
     #include "src/ext.h"
     """,
     extra_objects=['src/nearest_neighborhood.cu.o',
-                   '/usr/local/cuda-9.0/lib64/libcudart.so'],
+                   os.path.join(os.environ.get('CUDA_HOME'),'lib64/libcudart.so')],
     libraries=['stdc++']
 )