diff options
Diffstat (limited to 'torch')
| -rw-r--r-- | torch/backends/cudnn/__init__.py | 7 |
1 files changed, 6 insertions, 1 deletions
diff --git a/torch/backends/cudnn/__init__.py b/torch/backends/cudnn/__init__.py index 6c99146480..8d7542fdbc 100644 --- a/torch/backends/cudnn/__init__.py +++ b/torch/backends/cudnn/__init__.py @@ -74,12 +74,17 @@ CUDNN_TENSOR_TYPES = { } +def is_available(): + r"""Returns a bool indicating if CUDNN is currently available.""" + return torch._C.has_cudnn + + def is_acceptable(tensor): if not torch._C._get_cudnn_enabled(): return False if tensor.type() not in CUDNN_TENSOR_TYPES: return False - if not torch._C.has_cudnn: + if not is_available(): warnings.warn( "PyTorch was compiled without cuDNN support. To use cuDNN, rebuild " "PyTorch making sure the library is visible to the build system.") |