File size: 435 Bytes
bc1b8eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | from setuptools import setup
from torch.utils.cpp_extension import CppExtension, BuildExtension
setup(
name='sparse_linear',
ext_modules=[
CppExtension(
name='sparse_linear',
sources=['sparse_linear_ops.mm'],
extra_compile_args=['-ObjC++'],
extra_link_args=['-framework', 'Metal', '-framework', 'Foundation']
)
],
cmdclass={'build_ext': BuildExtension}
) |