Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
torch-cluster
Commits
a4b38f32
Commit
a4b38f32
authored
Nov 20, 2019
by
Dmitry Leonov
Browse files
Added flag for CPU install
parent
c7aeb8d2
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
8 additions
and
1 deletion
+8
-1
setup.py
setup.py
+8
-1
No files found.
setup.py
View file @
a4b38f32
from
setuptools
import
setup
,
find_packages
from
sys
import
argv
import
torch
from
torch.utils.cpp_extension
import
CppExtension
,
CUDAExtension
,
CUDA_HOME
...
...
@@ -22,7 +23,13 @@ ext_modules = [
]
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
+=
[
CUDAExtension
(
'torch_cluster.graclus_cuda'
,
[
'cuda/graclus.cpp'
,
'cuda/graclus_kernel.cu'
],
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment