naver/PUMP

1# Copyright 2022-present NAVER Corp.2# CC BY-NC-SA 4.03# Available only for non-commercial use4 5from setuptools import setup6from torch import cuda7from torch.utils.cpp_extension import BuildExtension, CUDAExtension8 9# if you want to compile for all possible CUDA architectures10all_cuda_archs = [] #cuda.get_gencode_flags().replace('compute=','arch=').split()11 12setup(13    name='cuda_deepm',14    ext_modules = [15        CUDAExtension(16                name = 'cuda_deepm',17                sources = ["func.cpp", "kernels.cu"],18                extra_compile_args = dict(nvcc=['-O2']+all_cuda_archs, cxx=['-O2'])19                )20    ],21    cmdclass = {22        'build_ext': BuildExtension23    })24 25