fix

setup.py

@@ -7,7 +7,6 @@ from sys import argv
 import torch
 from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
 
-
 # Windows users: Edit both of these to contain your VS include path, i.e.
 # cxx_extra_compile_args = ['-I{VISUAL_STUDIO_DIR}\\include']
 # nvcc_extra_compile_args = [..., '-I{VISUAL_STUDIO_DIR}\\include']
@@ -32,10 +31,9 @@ cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension}
 ext_modules = []
 exts = [e.split(osp.sep)[-1][:-4] for e in glob(osp.join('cpu', '*.cpp'))]
 ext_modules += [
-    CppExtension(
-        f'torch_scatter.{ext}_cpu', [f'cpu/{ext}.cpp'],
-        extra_compile_args=cxx_extra_compile_args,
-        extra_link_args=cxx_extra_link_args) for ext in exts
+    CppExtension(f'torch_scatter.{ext}_cpu', [f'cpu/{ext}.cpp'],
+                 extra_compile_args=cxx_extra_compile_args,
+                 extra_link_args=cxx_extra_link_args) for ext in exts
 ]
 
 if CUDA_HOME is not None and '--cpu' not in argv:
@@ -43,13 +41,13 @@ if CUDA_HOME is not None and '--cpu' not in argv:
     ext_modules += [
         CUDAExtension(
             f'torch_scatter.{ext}_cuda',
-            [f'cuda/{ext}.cpp', f'cuda/{ext}_kernel.cu'],
-            extra_compile_args={
+            [f'cuda/{ext}.cpp', f'cuda/{ext}_kernel.cu'], extra_compile_args={
                 'cxx': cxx_extra_compile_args,
                 'nvcc': nvcc_extra_compile_args,
-            },
-            extra_link_args=nvcc_extra_link_args) for ext in exts
+            }, extra_link_args=nvcc_extra_link_args) for ext in exts
     ]
+if '--cpu' in argv:
+    argv.remove('--cpu')
 
 __version__ = '1.5.0'
 url = 'https://github.com/rusty1s/pytorch_scatter'