from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension

setup(
    name='my_pybind_ops',            # 编译后的包名
    ext_modules=[
        CppExtension(
            name='my_pybind_ops',    # 扩展模块名（import 时用的名字）
            sources=['test_pybind_expand.cpp'], # C++ 源文件列表
        ),
    ],
    cmdclass={
        'build_ext': BuildExtension  # 使用 PyTorch 提供的构建类
    }
)