Skip to content

GitLab

Commit f6f147a7 authored by AntoinePrv's avatar AntoinePrv
Browse files

Weaker test for cuda extension.

parent 9bc8aa96
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
setup.py
View file @ f6f147a7
from setuptools import setup, find_packages from setuptools import setup, find_packages
import torch import torch
from torch.utils.cpp_extension import BuildExtension, CUDAExtension from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME
__version__ = '0.2.2' __version__ = '0.2.2'
url = 'https://github.com/rusty1s/pytorch_sparse' url = 'https://github.com/rusty1s/pytorch_sparse'
...@@ -11,7 +11,7 @@ tests_require = ['pytest', 'pytest-cov'] ...@@ -11,7 +11,7 @@ tests_require = ['pytest', 'pytest-cov']
ext_modules = [] ext_modules = []
cmdclass = {} cmdclass = {}
if torch.cuda.is_available(): if CUDA_HOME is not None:
ext_modules += [ ext_modules += [
CUDAExtension( CUDAExtension(
'spspmm_cuda', 'spspmm_cuda',
......
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