decagon-pytorch/src/torch_stablesort/setup.py

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']
            } ) ],
            cmdclass={'build_ext': cpp_extension.BuildExtension})