mirrors
/
mmcv
mirror of https://github.com/open-mmlab/mmcv.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix tensorrt test failed with pytorch 1.8+ (#1464) 

* fix tensorrt test failed with pytorch 1.8+

* add comment
pull/1579/head
Ming-Hsuan-Tu 2021-12-10 19:41:37 +08:00 committed by GitHub
parent 222f38075b
commit 48f44b697a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
setup.py
View File

@ -161,6 +161,9 @@ def get_extensions():
define_macros += [('MMCV_WITH_TRT', None)]
cuda_args = os.getenv('MMCV_CUDA_ARGS')
extra_compile_args['nvcc'] = [cuda_args] if cuda_args else []
# prevent cub/thrust conflict with other python library
# More context See issues #1454
extra_compile_args['nvcc'] += ['-Xcompiler=-fno-gnu-unique']
library_dirs += library_paths(cuda=True)
from setuptools import Extension