Update setup.py

PyTorch/setup.py

@@ -13,13 +13,21 @@ torch_dir = os.path.dirname(torch.__file__)
 print('Building SCN module')
 if torch.cuda.is_available():
     r = os.system(
-        'cd sparseconvnet/SCN; nvcc init.cu -c -o init.cu.o -ccbin /usr/bin/cc -m64 --std c++11 -Xcompiler ,\"-fopenmp\",\"-fPIC\",\"-O3\",\"-DNDEBUG\" -gencode arch=compute_52,code=sm_52 -gencode arch=compute_30,code=sm_30 -DNVCC -I/usr/local/cuda/include -I' +
-        torch_dir +
-        '/lib/include -I' +
-        torch_dir +
-        '/lib/include/TH -I' +
-        torch_dir +
-        '/lib/include/THC -I.')
+        'cd sparseconvnet/SCN; nvcc init.cu -c -o init.cu.o -ccbin /usr/bin/cc'
+        + ' -m64 --std c++11 -Xcompiler '
+        + ',\"-fopenmp\",\"-fPIC\",\"-O3\",\"-DNDEBUG\" '
+        + '-gencode arch=compute_62,code=sm_62 '
+        + '-gencode arch=compute_61,code=sm_61 '
+        + '-gencode arch=compute_60,code=sm_60 '
+        + '-gencode arch=compute_52,code=sm_52 '
+        + '-gencode arch=compute_50,code=sm_50 '
+        + '-gencode arch=compute_30,code=sm_30 '
+        + '-DNVCC '
+        + '-I/usr/local/cuda/include '
+        + '-I' + torch_dir + '/lib/include '
+        + '-I' + torch_dir + '/lib/include/TH '
+        + '-I' + torch_dir + '/lib/include/THC '
+        + '-I.')
     assert r == 0
     ffi = create_extension(
         'sparseconvnet.SCN',
@@ -60,5 +68,10 @@ setup(
     author='Facebook AI Research',
     author_email='benjamingraham@fb.com',
     url='https://github.com/facebookresearch/SparseConvNet',
+    package_data={
+        'sparseconvnet': ['SCN/_SCN.so'],
+    },
     packages=find_packages(),
+    # Since the package includes a shared object, this is not zip-safe.
+    zip_safe=False,
 )