|
1234567891011121314 |
- from setuptools import setup, Extension
- from torch.utils import cpp_extension
-
- setup(name='torch_stablesort',
- py_modules=['torch_stablesort'],
- ext_modules=[ cpp_extension.CUDAExtension( 'torch_stablesort_cpp',
- ['torch_stablesort.cpp', 'torch_stablesort_cpu.cpp', 'torch_stablesort_cuda.cu'],
- extra_compile_args={
- 'cxx': ['-fopenmp', '-ggdb', '-std=c++1z'],
- 'nvcc': [ '-I/pstore/home/adaszews/scratch/thrust',
- '-ccbin', '/pstore/data/data_science/app/modules/anaconda3-2020.07/bin/x86_64-conda_cos6-linux-gnu-gcc',
- '-std=c++14', '--expt-extended-lambda', '-O99']
- } ) ],
- cmdclass={'build_ext': cpp_extension.BuildExtension})
|