Skip to content

GitLab

Commit 197bcc48 authored by Michael Carilli's avatar Michael Carilli
Browse files

Warning instead of error if nvcc is not found

parent b56a2088
Show whitespace changes
Inline Side-by-side
Showing with 24 additions and 20 deletions
setup.py
View file @ 197bcc48
...@@ -32,6 +32,10 @@ if "--cpp_ext" in sys.argv: ...@@ -32,6 +32,10 @@ if "--cpp_ext" in sys.argv:
if "--cuda_ext" in sys.argv: if "--cuda_ext" in sys.argv:
from torch.utils.cpp_extension import CUDAExtension from torch.utils.cpp_extension import CUDAExtension
sys.argv.remove("--cuda_ext") sys.argv.remove("--cuda_ext")
if torch.utils.cpp_extension.CUDA_HOME is None:
print("Warning: nvcc is not available. Ignoring --cuda-ext")
else:
ext_modules.append( ext_modules.append(
CUDAExtension(name='fused_adam_cuda', CUDAExtension(name='fused_adam_cuda',
sources=['apex/optimizers/csrc/fused_adam_cuda.cpp', sources=['apex/optimizers/csrc/fused_adam_cuda.cpp',
......
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