Merge pull request #15 from AntoinePrv/cuda_no_cuda

weaker test for building cuda extension.

Merge pull request #15 from AntoinePrv/cuda_no_cuda
weaker test for building cuda extension.
15afee0d · Matthias Fey · GitHub · 77b4b50b · ce46b2ce · 15afee0d
Commit 15afee0d authored Nov 23, 2018 by Matthias Fey Committed by GitHub Nov 23, 2018
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'])