Python is_tensor_core_availableの例

コード例 #1

def should_use_cudnn_tensor_core(dtype):
    """Determines if Tensor Core should be used.

    Args:
        dtype (numpy.dtype): data type of input tensor.

    Returns:
        bool: ``True`` if Tensor Core should be used.
    """

    use_cudnn_tensor_core = config.use_cudnn_tensor_core
    if use_cudnn_tensor_core not in _tensor_core_flag:
        raise ValueError('invalid use_cudnn_tensor_core configuration: %s '
                         '(must be either of "always", "auto", or "never")' %
                         repr(use_cudnn_tensor_core))
    use_tensor_core = _tensor_core_flag[use_cudnn_tensor_core]
    if use_tensor_core is None:
        use_tensor_core = cudnn.is_tensor_core_available(dtype)
    return use_tensor_core

コード例 #2

ファイル: cuda.py プロジェクト: jnishi/chainer

def should_use_cudnn_tensor_core(dtype):
    """Determines if Tensor Core should be used.

    Args:
        dtype (numpy.dtype): data type of input tensor.

    Returns:
        bool: ``True`` if Tensor Core should be used.
    """

    use_cudnn_tensor_core = config.use_cudnn_tensor_core
    if use_cudnn_tensor_core not in _tensor_core_flag:
        raise ValueError('invalid use_cudnn_tensor_core configuration: %s '
                         '(must be either of "always", "auto", or "never")' %
                         repr(use_cudnn_tensor_core))
    use_tensor_core = _tensor_core_flag[use_cudnn_tensor_core]
    if use_tensor_core is None:
        use_tensor_core = cudnn.is_tensor_core_available(dtype)
    return use_tensor_core