diff options
| author | Peter Goldsborough <peter@goldsborough.me> | 2018-06-20 14:18:03 -0700 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2018-06-20 14:18:03 -0700 |
| commit | 0acddd6cee32bc7c3715bc8b93d0a33ef19064b1 (patch) | |
| tree | dddf08ca89ddc64a1870038c2ae52d68d62f4e6c /torch | |
| parent | 85468155cee8a08fd7437f8addfecb18f310c6a8 (diff) | |
| download | pytorch-0acddd6cee32bc7c3715bc8b93d0a33ef19064b1.tar.gz pytorch-0acddd6cee32bc7c3715bc8b93d0a33ef19064b1.tar.bz2 pytorch-0acddd6cee32bc7c3715bc8b93d0a33ef19064b1.zip | |
Add torch.cuda.cudnn_is_available (#8703)
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.") |