Commit fff1994e authored by Rick Ho's avatar Rick Ho
Browse files

optional cuda helper path by environ var

parent b83ac1a5
from setuptools import setup from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension from torch.utils.cpp_extension import BuildExtension, CUDAExtension
import os
CUDA_HELPER = os.environ.get('CUDA_HELPER', '/usr/local/cuda/samples/common/inc')
setup( setup(
name='moe_cuda', name='moe_cuda',
...@@ -10,8 +13,8 @@ setup( ...@@ -10,8 +13,8 @@ setup(
'moe.cpp', 'moe.cpp',
'moe_cuda_kernel.cu', 'moe_cuda_kernel.cu',
], ],
extra_compile_args={'cxx': ['-I/usr/local/cuda/samples/common/inc'], extra_compile_args={'cxx': ['-I{}'.format(CUDA_HELPER)],
'nvcc': ['-I/usr/local/cuda/samples/common/inc']} 'nvcc': ['-I{}'.format(CUDA_HELPER)]}
) )
], ],
cmdclass={ cmdclass={
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment