diff options
| author | Sam Gross <colesbury@gmail.com> | 2016-09-28 19:29:53 -0400 |
|---|---|---|
| committer | Adam Paszke <adam.paszke@gmail.com> | 2016-09-28 19:29:53 -0400 |
| commit | cb5d4e836fa015c0364bca962292a0de6b78a8e5 (patch) | |
| tree | e231b5a3ee6d1ba4e3fcdc104db01a7024cca444 /torch/backends | |
| parent | 412019dbe4998bc56aaa03b30280e65b23a2fe35 (diff) | |
| download | pytorch-cb5d4e836fa015c0364bca962292a0de6b78a8e5.tar.gz pytorch-cb5d4e836fa015c0364bca962292a0de6b78a8e5.tar.bz2 pytorch-cb5d4e836fa015c0364bca962292a0de6b78a8e5.zip | |
Lazy load CUDA and THNN modules (#64)
Diffstat (limited to 'torch/backends')
| -rw-r--r-- | torch/backends/cudnn/__init__.py | 14 |
1 files changed, 5 insertions, 9 deletions
diff --git a/torch/backends/cudnn/__init__.py b/torch/backends/cudnn/__init__.py index 3fa5c46262..938d1b4fb4 100644 --- a/torch/backends/cudnn/__init__.py +++ b/torch/backends/cudnn/__init__.py @@ -1,5 +1,6 @@ import ctypes import warnings +import torch.cuda lib = None libname = 'libcudnn.so.5.1.3' @@ -12,24 +13,19 @@ def _loadlib(): def is_acceptable(tensor): + if not (isinstance(tensor, torch.cuda.HalfTensor) or + isinstance(tensor, torch.cuda.FloatTensor) or + isinstance(tensor, torch.cuda.DoubleTensor)): + return False if lib is None: try: - import torch.cuda - except ImportError: - return False - try: _loadlib() except Exception: warnings.warn('cuDNN library not found. Check your LD_LIBRARY_PATH') return False - if not (isinstance(tensor, torch.cuda.HalfTensor) or - isinstance(tensor, torch.cuda.FloatTensor) or - isinstance(tensor, torch.cuda.DoubleTensor)): - return False return True - _handles = {} benchmark = False |