Hi, everyone. I want a rank layer and then use backpropagation transfer its gradient to its original position, and I have written a custom function. When I test it, the forward propagation works well, and I can get expected results. However,  I can not get the result of gradient. I’m really curious about where I have ignored.
This is my custom function:
class RankingFunc(Function): 
    @staticmethod
    def forward(ctx, x, ranking_size): 
        p, c, h, w = x.size()
        y = torch.empty(x.size())
        ind = torch.arange(0, h*w).view(h, w).unsqueeze(0).repeat(p, c, 1, 1)
        ctx.ranking_size = ranking_size
        
        h_num = h // ranking_size
        w_num = w // ranking_size
        
        for i in range(h_num):
            for j in range(w_num):
                temp1 = x[:, :, i * ranking_size : (i + 1) * ranking_size, j * ranking_size : (j + 1) * ranking_size].clone()
                temp2 = ind[:, :, i * ranking_size : (i + 1) * ranking_size, j * ranking_size : (j + 1) * ranking_size].clone()
                value, index = torch.sort(temp1.view(p, c, ranking_size ** 2))
                temp1 = value.view(p, c, ranking_size, ranking_size)
                y[:, :, i * ranking_size : (i + 1) * ranking_size, j * ranking_size : (j + 1) * ranking_size] = temp1
                temp2 = temp2.view(p, c, ranking_size ** 2).view(-1)
                index = index.view(-1)
                temp2 = temp2[index].view(p, c, ranking_size, ranking_size)
                ind[:, :, i * ranking_size : (i + 1) * ranking_size, j * ranking_size : (j + 1) * ranking_size] = temp2
        ctx.save_for_backward(ind)
        return y
    
    @staticmethod
    def backward(ctx, grad_y):
        index = ctx.saved_variables[0]
        p, c, h, w = grad_y.size()
        grad_y = grad_y.view(p, c, -1).contiguous()
        _, index = index.view(p, c, -1).sort()
        grad_y = grad_y.view(-1)[index.view(-1)]
        grad_y = grad_y.view(p, c, h, w)
        return grad_y, None

And this is my test code:
import torch
import RankingFunc

x = torch.rand((2, 2, 4, 4), requires_grad = True) * 100
y = RankingFunc.apply(x, 3)
z = y / 2

z_grad = torch.rand(2, 2, 4, 4) * 100
z.backward(z_grad)

There is just no output when I print(x.grad). Does anybody have any ideas? Thanks in advance!