Hi Thomas, I still have a problem on input.grad. The output is always None, do you know what the wrong is?
here is my class py file:
class GlobalContrastFunction(Function):

    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x) # save x for backward
        return global_contrast.forward(x)

    @staticmethod
    def backward(ctx, grad):
        x = ctx.saved_variables # load
        dx = global_contrast.backward(grad.contiguous(), *x) # get gradient 
        print(dx.data[0,0,:4,:4])  # print out shows that dx has correct value!
        return dx


class GlobalContrast(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return GlobalContrastFunction.apply(x)

while in the loop of main py file:
    globalContrast = GlobalContrast().cuda()
    x = torch.rand((B, C, W, H), requires_grad=True).cuda()

    for _ in tqdm(range(options.runs)):

        if options.example == 'cuda':
            globalContrast.zero_grad()

        c = globalContrast(x) # forward

        grad = torch.sum(c)
        grad.backward(retain_graph=True) # set retain_graph=True here to save gradients for each node

        print(c.grad) # Here is None!
        print(x.grad) # Here is also None!