From 3e4a963d8fe304e500b5a279b2e0b249c26de390 Mon Sep 17 00:00:00 2001
From: AntoinePrv <antoine.prouvost.2013_dev@polytechnique.org>
Date: Thu, 22 Nov 2018 15:52:34 -0500
Subject: [PATCH] weaker test for building cuda extension.

---
 setup.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index 652ac40c..7919433f 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,6 @@
 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 @@
 ]
 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'])