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.CppExtension('torch_stablesort_cpp',
        ['torch_stablesort.cpp'],
        extra_compile_args=['-fopenmp', '-ggdb'])],
            cmdclass={'build_ext': cpp_extension.BuildExtension})