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

Warning instead of error if nvcc is not found

parent b56a2088
...@@ -19,7 +19,7 @@ Installation requires CUDA 9 or later, PyTorch 0.4 or later, and Python 3. Insta ...@@ -19,7 +19,7 @@ Installation requires CUDA 9 or later, PyTorch 0.4 or later, and Python 3. Insta
cd apex cd apex
python setup.py install [--cuda_ext] [--cpp_ext] python setup.py install [--cuda_ext] [--cpp_ext]
.. toctree:: .. toctree::
:maxdepth: 1 :maxdepth: 1
:caption: AMP: Automatic Mixed Precision :caption: AMP: Automatic Mixed Precision
......
...@@ -32,25 +32,29 @@ if "--cpp_ext" in sys.argv: ...@@ -32,25 +32,29 @@ 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")
ext_modules.append(
CUDAExtension(name='fused_adam_cuda', if torch.utils.cpp_extension.CUDA_HOME is None:
sources=['apex/optimizers/csrc/fused_adam_cuda.cpp', print("Warning: nvcc is not available. Ignoring --cuda-ext")
'apex/optimizers/csrc/fused_adam_cuda_kernel.cu'], else:
extra_compile_args={'cxx': ['-O3',], ext_modules.append(
'nvcc':['-O3', CUDAExtension(name='fused_adam_cuda',
'--use_fast_math']})) sources=['apex/optimizers/csrc/fused_adam_cuda.cpp',
ext_modules.append( 'apex/optimizers/csrc/fused_adam_cuda_kernel.cu'],
CUDAExtension(name='syncbn', extra_compile_args={'cxx': ['-O3',],
sources=['csrc/syncbn.cpp', 'nvcc':['-O3',
'csrc/welford.cu'])) '--use_fast_math']}))
ext_modules.append( ext_modules.append(
CUDAExtension(name='fused_layer_norm_cuda', CUDAExtension(name='syncbn',
sources=['apex/normalization/csrc/layer_norm_cuda.cpp', sources=['csrc/syncbn.cpp',
'apex/normalization/csrc/layer_norm_cuda_kernel.cu'], 'csrc/welford.cu']))
extra_compile_args={'cxx': ['-O3',], ext_modules.append(
'nvcc':['-maxrregcount=50', CUDAExtension(name='fused_layer_norm_cuda',
'-O3', sources=['apex/normalization/csrc/layer_norm_cuda.cpp',
'--use_fast_math']})) 'apex/normalization/csrc/layer_norm_cuda_kernel.cu'],
extra_compile_args={'cxx': ['-O3',],
'nvcc':['-maxrregcount=50',
'-O3',
'--use_fast_math']}))
setup( setup(
name='apex', name='apex',
......
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