"git@developer.sourcefind.cn:OpenDAS/apex.git" did not exist on "8124df1369ca927305ea4de0e8ffff9d083e6eb2"
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,10 +13,10 @@ setup( ...@@ -10,10 +13,10 @@ 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={
'build_ext': BuildExtension 'build_ext': BuildExtension
}) })
\ No newline at end of file
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