Python torch.backends.cudnn 模块，enabled() 实例源码

def _enforce_cudnn(input):
        if not cudnn.enabled:
            raise RuntimeError("AffineGridGenerator needs CuDNN for "
                               "processing CUDA inputs, but CuDNN is not enabled")
        assert cudnn.is_acceptable(input)

def use_cudnn(should_use):
    orig = torch.backends.cudnn.enabled
    torch.backends.cudnn.enabled = should_use
    try:
        yield
    finally:
        torch.backends.cudnn.enabled = orig

def test_Conv2d_inconsistent_types_on_GPU_without_cudnn(self):
        inputs = Variable(torch.randn(4, 1, 7, 7).float().cuda())
        weights = Variable(torch.randn(1, 1, 3, 3).double().cuda())
        bias = Variable(torch.randn(1).double().cuda())

        torch.backends.cudnn.enabled = False
        # inconsistent types should raise an exception
        self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights))
        self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights.float(), bias))

        # but it should work with the same type
        nn.functional.conv2d(inputs.float(), weights.float(), bias.float())

def test_Conv2d_inconsistent_types_on_GPU_with_cudnn(self):
        inputs = Variable(torch.randn(4, 1, 7, 7).float().cuda())
        weights = Variable(torch.randn(1, 1, 3, 3).double().cuda())
        bias = Variable(torch.randn(1).double().cuda())

        torch.backends.cudnn.enabled = True
        # inconsistent types should raise an exception
        self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights))
        self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights.float(), bias))

        # but it should work with the same type
        nn.functional.conv2d(inputs.float(), weights.float(), bias.float())

def test_rnn_retain_variables_cuda(self):
        try:
            torch.backends.cudnn.enabled = False
            self._test_rnn_retain_variables(torch.cuda.FloatTensor)
        finally:
            torch.backends.cudnn.enabled = True
        self._test_rnn_retain_variables(torch.cuda.FloatTensor)

def affine_grid_generator(theta, size):
    if theta.data.is_cuda:
        if not cudnn.enabled:
            raise RuntimeError("AffineGridGenerator needs CuDNN for "
                               "processing CUDA inputs, but CuDNN is not enabled")
        if not cudnn.is_acceptable(theta.data):
            raise RuntimeError("AffineGridGenerator generator theta not acceptable for CuDNN")
        N, C, H, W = size
        return torch._C._VariableBase.cudnn_affine_grid_generator(theta, N, C, H, W)
    else:
        return AffineGridGenerator.apply(theta, size)


# TODO: Port these completely into C++

def _enforce_cudnn(input):
        if not cudnn.enabled:
            raise RuntimeError("AffineGridGenerator needs CuDNN for "
                               "processing CUDA inputs, but CuDNN is not enabled")
        assert cudnn.is_acceptable(input)