saymrwulf/pytorch
1
0
Fork 0
mirror of https://github.com/saymrwulf/pytorch.git synced 2026-05-14 20:57:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add torch.cuda.cudnn_is_available (#8703)

Browse source
This commit is contained in:
Peter Goldsborough 2018-06-20 14:18:03 -07:00 committed by GitHub
parent 85468155ce
commit 0acddd6cee
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
torch/backends/cudnn/__init__.py
View file

@ -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.")