Arnaudding001 commited on
Commit
18a919c
1 Parent(s): 511a2cd

Create raft_alt_cuda_corr_setup.py

Browse files
Files changed (1) hide show
  1. raft_alt_cuda_corr_setup.py +14 -0
raft_alt_cuda_corr_setup.py ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from setuptools import setup
2
+ from torch.utils.cpp_extension import BuildExtension, CUDAExtension
3
+
4
+
5
+ setup(
6
+ name='correlation',
7
+ ext_modules=[
8
+ CUDAExtension('alt_cuda_corr',
9
+ sources=['correlation.cpp', 'correlation_kernel.cu'],
10
+ extra_compile_args={'cxx': [], 'nvcc': ['-O3']}),
11
+ ],
12
+ cmdclass={
13
+ 'build_ext': BuildExtension
14
+ })