Skip to content

GitLab

Commit 15afee0d authored by Matthias Fey's avatar Matthias Fey Committed by GitHub
Browse files

Merge pull request #15 from AntoinePrv/cuda_no_cuda 

weaker test for building cuda extension.
parents 77b4b50b ce46b2ce
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
setup.py
View file @ 15afee0d
from setuptools import setup, find_packages from setuptools import setup, find_packages
import torch import torch
from torch.utils.cpp_extension import CppExtension, CUDAExtension from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
ext_modules = [ ext_modules = [
CppExtension( CppExtension(
...@@ -9,7 +9,7 @@ ext_modules = [ ...@@ -9,7 +9,7 @@ ext_modules = [
] ]
cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension} cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension}
if torch.cuda.is_available(): if CUDA_HOME is not None:
ext_modules += [ ext_modules += [
CUDAExtension('scatter_cuda', CUDAExtension('scatter_cuda',
['cuda/scatter.cpp', 'cuda/scatter_kernel.cu']) ['cuda/scatter.cpp', 'cuda/scatter_kernel.cu'])
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment