fix: setup.py TORCH_CUDA_ARCH_LIST (#62)
This commit is contained in:
2
setup.py
2
setup.py
@@ -70,7 +70,7 @@ def get_extensions():
|
|||||||
extra_compile_args = {"cxx": []}
|
extra_compile_args = {"cxx": []}
|
||||||
define_macros = []
|
define_macros = []
|
||||||
|
|
||||||
if torch.cuda.is_available() and CUDA_HOME is not None:
|
if CUDA_HOME is not None and (torch.cuda.is_available() or "TORCH_CUDA_ARCH_LIST" in os.environ):
|
||||||
print("Compiling with CUDA")
|
print("Compiling with CUDA")
|
||||||
extension = CUDAExtension
|
extension = CUDAExtension
|
||||||
sources += source_cuda
|
sources += source_cuda
|
||||||
|
Reference in New Issue
Block a user