Commit a4b38f32 authored by Dmitry Leonov's avatar Dmitry Leonov
Browse files

Added flag for CPU install

parent c7aeb8d2
from setuptools import setup, find_packages from setuptools import setup, find_packages
from sys import argv
import torch import torch
from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
...@@ -22,7 +23,13 @@ ext_modules = [ ...@@ -22,7 +23,13 @@ ext_modules = [
] ]
cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension} cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension}
if CUDA_HOME is not None: GPU = True
for arg in argv:
if arg == '--cpu':
GPU = False
argv.remove(arg)
if CUDA_HOME is not None and GPU:
ext_modules += [ ext_modules += [
CUDAExtension('torch_cluster.graclus_cuda', CUDAExtension('torch_cluster.graclus_cuda',
['cuda/graclus.cpp', 'cuda/graclus_kernel.cu'], ['cuda/graclus.cpp', 'cuda/graclus_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