from setuptools import find_packages, setup
from torch.utils.cpp_extension import CUDAExtension, BuildExtension

def get_extensions():
 
    return extensions


if __name__ == "__main__":
    extensions = [
        CUDAExtension(
            "broadcast",
            sources=[
                "Cluster-ViT/models/extensions/broadcast.cu"
            ],
            extra_compile_args=["-arch=compute_50"]
        ),
        CUDAExtension(
            "weighted_sum",
            sources=[
                "Cluster-ViT/models/extensions/weighted_sum.cu"
            ],
            extra_compile_args=["-arch=compute_50"]
        )
    ]

    setup(
        name="clutering-Transformer",
        packages=find_packages(),
        ext_modules=extensions,
        cmdclass={"build_ext": BuildExtension},
        install_requires=["torch"]
    )