I use package Numba to write the forward and backward cuda operation, and use it as a module. When I use it, the input tensors, with grad required, go through the modules. The example can be described as:
def forward_cudakernel():
        ...   # write by NUMBA

def backward_cudakernel():
        ...  # write by NUMBA

def forward_cuda(x)
    with torch.no_grad():
        x = numba.cuda.as_cuda_array(x)
        forward_cudakernel[blocks, threads](x)

def backward_cuda(x)
    with torch.no_grad():
        x = numba.cuda.as_cuda_array(x)
        backward_cudakernel[blocks, threads](x)

class BaseModule(Function):
    @staticmethod
    def forward():
        ...
        return forward_cuda()
    @staticmethod
    def backward():
        ...
        return grads

class Module(nn.Module):
    def forward(feature):
        fun = BaseModule.apply
        return fun(feature)

>>> features = torch.rand(2, 256, 20, 20).cuda()
>>> features.requires_grad_()
>>> net = Module()
>>> out = net(features)
>>> out.mean.backward()

Where “out” and “feature” needs grads so that the module can run. However, the attr “cuda_array_interface”（cuda.as_cuda_array related） is not available when a tensor is with gradient. It seems incompatible. So, How can I use “numba” and “pytorch” together? Thanks!