Unverified Commit b294221d authored by Rick Ho's avatar Rick Ho Committed by GitHub
Browse files

update setup.py to be compatible with torch v1.10.0 (#88)

parent acf8bec8
...@@ -22,6 +22,7 @@ if torch.__version__ >= '1.5': ...@@ -22,6 +22,7 @@ if torch.__version__ >= '1.5':
if os.environ.get('USE_NCCL', '1') == '1': if os.environ.get('USE_NCCL', '1') == '1':
cxx_flags.append('-DFMOE_USE_NCCL') cxx_flags.append('-DFMOE_USE_NCCL')
cxx_flags.append('-DUSE_C10D_NCCL')
if is_rocm_pytorch: if is_rocm_pytorch:
ext_libs.append('rccl') ext_libs.append('rccl')
else: else:
......
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