Commit 02ead75d authored by Rick Ho's avatar Rick Ho
Browse files

add nccl to lib list

parent da11cb76
......@@ -7,8 +7,10 @@ CUDA_HELPER = os.environ.get('CUDA_HELPER', '/usr/local/cuda/samples/common/inc'
cxx_flags = [
'-I{}'.format(CUDA_HELPER)
]
ext_libs = []
if os.environ.get('USE_NCCL', '0') == '1':
cxx_flags.append('-DMOE_USE_NCCL')
ext_libs.append('nccl')
if __name__ == '__main__':
......@@ -34,7 +36,8 @@ if __name__ == '__main__':
extra_compile_args={
'cxx': cxx_flags,
'nvcc': cxx_flags
}
},
libraries=ext_libs
)
],
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