setup.py 343 B

123456789101112
  1. from setuptools import setup
  2. from torch.utils.cpp_extension import BuildExtension, CUDAExtension
  3. setup(name='swin_window_process',
  4. ext_modules=[
  5. CUDAExtension('swin_window_process', [
  6. 'swin_window_process.cpp',
  7. 'swin_window_process_kernel.cu',
  8. ])
  9. ],
  10. cmdclass={'build_ext': BuildExtension})