weaker test for building cuda extension.

ce46b2ce · AntoinePrv · 77b4b50b · ce46b2ce
Commit ce46b2ce authored Nov 22, 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 CppExtension, CUDAExtension
+from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME

 ext_modules = [
    CppExtension(
@@ -9,7 +9,7 @@ ext_modules = [
 ]
 cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension}

-if torch.cuda.is_available():
+if CUDA_HOME is not None:
    ext_modules += [
        CUDAExtension('scatter_cuda',
                      ['cuda/scatter.cpp', 'cuda/scatter_kernel.cu'])