When I write cpp extension for custom cudnn convolution, I use nn.autograd and nn.Module wrap my cpp extension.
autograd wraper code in Cudnn_conv2d_func.py file like this:
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Function
import math
import cudnn_conv2d


class Cudnn_conv2d_func(Function):
    @staticmethod
    def forward(ctx, input, weight, bias=None, stride=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1, benchmark=True, deterministic=False):
        ctx.stride = stride
        ctx.padding= padding
        ctx.dilation = dilation
        ctx.groups = groups
        ctx.benchmark = benchmark
        ctx.deterministic = deterministic
        ctx.save_for_backward(input, weight)
        stride_width, stride_height = stride[0], stride[1]
        padding_width, padding_height = padding[0], padding[1]
        dilation_width, dilation_height = dilation[0], dilation[1]

        return cudnn_conv2d.conv2d(input, weight, bias, stride_width, stride_height, padding_width,
                                   padding_height, dilation_width, dilation_height, groups, benchmark, deterministic)


    @staticmethod
    def backward(ctx, grad_output):
        input, weight = ctx.saved_variables

        print(ctx.needs_input_grad)

        grad_input, grad_weight, grad_bias = cudnn_conv2d.conv2d_backward(
                                                            input, grad_output, weight, 
                                                            ctx.stride[0], ctx.stride[1],
                                                            ctx.padding[0], ctx.padding[1],
                                                            ctx.dilation[0], ctx.dilation[1],
                                                            ctx.groups, ctx.benchmark, ctx.deterministic,
                                                            ctx.needs_input_grad[:3])
        
        return grad_input, grad_weight, None, None, None, None, None, None


nn.Module wraper code in Cudnn_conv2d.py file like this:
import torch
from torch.nn.parameter import Parameter
from torch.nn import init
from torch.nn.modules import Module
from torch.nn.modules.utils import _single, _pair, _triple
from torch._jit_internal import List
from .Cudnn_conv2d_func import Cudnn_conv2d_func
import math


class _ConvNd(Module):

    __constants__ = ['stride', 'padding', 'dilation', 'groups', 'bias',
                     'padding_mode', 'output_padding', 'in_channels',
                     'out_channels', 'kernel_size']

    def __init__(self, in_channels, out_channels, kernel_size, stride,
                 padding, dilation, transposed, output_padding,
                 groups, bias, padding_mode):
        super(_ConvNd, self).__init__()
        if in_channels % groups != 0:
            raise ValueError('in_channels must be divisible by groups')
        if out_channels % groups != 0:
            raise ValueError('out_channels must be divisible by groups')
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.transposed = transposed
        self.output_padding = output_padding
        self.groups = groups
        self.padding_mode = padding_mode
        if transposed:
            self.weight = Parameter(torch.Tensor(
                in_channels, out_channels // groups, *kernel_size))
        else:
            self.weight = Parameter(torch.Tensor(
                out_channels, in_channels // groups, *kernel_size))
        if bias:
            self.bias = Parameter(torch.Tensor(out_channels))
        else:
            # self.register_parameter('bias', None)
            self.bias = Parameter(torch.zeros(out_channels))
        self.reset_parameters()

    def reset_parameters(self):
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in)
            init.uniform_(self.bias, -bound, bound)

    def extra_repr(self):
        s = ('{in_channels}, {out_channels}, kernel_size={kernel_size}'
             ', stride={stride}')
        if self.padding != (0,) * len(self.padding):
            s += ', padding={padding}'
        if self.dilation != (1,) * len(self.dilation):
            s += ', dilation={dilation}'
        if self.output_padding != (0,) * len(self.output_padding):
            s += ', output_padding={output_padding}'
        if self.groups != 1:
            s += ', groups={groups}'
        if self.bias is None:
            s += ', bias=False'
        if self.padding_mode != 'zeros':
            s += ', padding_mode={padding_mode}'
        return s.format(**self.__dict__)

    def __setstate__(self, state):
        super(_ConvNd, self).__setstate__(state)
        if not hasattr(self, 'padding_mode'):
            self.padding_mode = 'zeros'


class Cudnn_conv2d(_ConvNd):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 padding=0, dilation=1, groups=1,
                 bias=True, padding_mode='zeros'):
        kernel_size = _pair(kernel_size)
        stride = _pair(stride)
        padding = _pair(padding)
        dilation = _pair(dilation)
        super(Cudnn_conv2d, self).__init__(
            in_channels, out_channels, kernel_size, stride, padding, dilation,
            False, _pair(0), groups, bias, padding_mode)

    def conv2d_forward(self, input, weight):
        if self.padding_mode == 'circular':
            expanded_padding = ((self.padding[1] + 1) // 2, self.padding[1] // 2,
                                (self.padding[0] + 1) // 2, self.padding[0] // 2)
            return Cudnn_conv2d_func.apply(F.pad(input, expanded_padding, mode='circular'), weight, self.bias, self.stride, _pair(0), self.dilation, self.groups)

        return Cudnn_conv2d_func.apply(input, weight, self.bias, self.stride, self.padding, self.dilation, self.groups)

    def forward(self, input):
        return self.conv2d_forward(input, self.weight)



if __name__=='__main__':
    conv = Cudnn_conv2d(3, 3, 3).cuda()
    print(conv)
    x = torch.rand([128, 3, 32, 32],requires_grad=True, device='cuda')
    y = conv(x)
    y = y.sum()
    print(y)
    y.backward()

When I run Cudnn_conv2d.py, I can get right result. But When I use Cudnn_conv2d instead nn.Conv2d  in net. I meet Error like this;
ctx.needs_input_grad[:3])
RuntimeError: Expected tensor's dynamic type to be Variable, not Tensor

I am confusing. Net code look like this:
import torch
import torch.nn as nn
import torch.nn.functional as F
from .Cudnn_conv2d import Cudnn_conv2d


class TestNet1(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1=Cudnn_conv2d(3,32,5,bias=False)
        self.conv2=Cudnn_conv2d(32,32,3,bias=False)
        self.conv3=Cudnn_conv2d(32,32,3,bias=False)
        self.bn = nn.BatchNorm2d(32)
        self.fc1=nn.Linear(16*32,10)

    def forward(self,x):
        x=self.conv1(x) #28
        x=F.relu(x)
        x=F.max_pool2d(x,2,2) #14
        x=self.conv2(x) # 12
        x = self.bn(x)
        x=F.relu(x)
        x=F.max_pool2d(x,2,2) # 6
        x=self.conv3(x) # 4
        x=x.view(x.size(0),-1)
        x=self.fc1(x)
        return x

dataset is cifar10.
Does anybody know where to look to debug?
Thanks!