Weaker test for cuda extension.

f6f147a7 · AntoinePrv · 9bc8aa96 · f6f147a7
Commit f6f147a7 authored Nov 23, 2018 by AntoinePrv
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

setup.py setup.py +2 -2

No files found.
--- a/setup.py
+++ b/setup.py
 from setuptools import setup, find_packages
 import torch
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME

 __version__ = '0.2.2'
 url = 'https://github.com/rusty1s/pytorch_sparse'
@@ -11,7 +11,7 @@ tests_require = ['pytest', 'pytest-cov']
 ext_modules = []
 cmdclass = {}

-if torch.cuda.is_available():
+if CUDA_HOME is not None:
    ext_modules += [
        CUDAExtension(
            'spspmm_cuda',