"src/git@developer.sourcefind.cn:renzhc/diffusers_dcu.git" did not exist on "05b706c003666709471c62bd44b8f40190506000"
Commit 61b9c55c authored by Dmitry Leonov's avatar Dmitry Leonov
Browse files

Added flag for CPU install

parent e8aa000a
import platform import platform
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
...@@ -16,7 +17,13 @@ ext_modules = [ ...@@ -16,7 +17,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:
if platform.system() == 'Windows': if platform.system() == 'Windows':
extra_link_args = ['cusparse.lib'] extra_link_args = ['cusparse.lib']
else: else:
......
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