setup.py

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

setup(
    name="lagomorph",
    author='Jacob Hinkle',
    author_email='hinklejd@ornl.gov',
    url='https://github.com/jacobhinkle/lagomorph',
    keywords='medical-imaging computation-anatomy image-registration computer-vision pytorch cuda',
    packages=['lagomorph'],
    python_requires=">=3.6",
    use_scm_version=True,
    setup_requires=['pytest-runner','setuptools_scm'],
    install_requires=['torch>=1.0','numpy','h5py'],
    tests_require=['pytest'],
    cmdclass={'build_ext': BuildExtension},
    entry_points={'console_scripts':['lagomorph=lagomorph.__main__:main']},
    ext_modules=[CUDAExtension('lagomorph_ext', [
            'lagomorph/extension/cpu/affine.cpp',
           'lagomorph/extension/cuda/affine.cu',
            'lagomorph/extension/cuda/diff.cu',
            'lagomorph/extension/cuda/interp.cu',
            'lagomorph/extension/cuda/metric.cu',
            'lagomorph/extension/extension.cpp'
        ],
        include_dirs=[os.path.abspath("lagomorph/extension/include")])]
)