关于torch.topK取得的元素可导性问题

先上结论是可导的，但是可导的也只有被筛选出来 那一部分

import torch

import torch
import torch.nn as nn


class g(nn.Module):
    def __init__(self):
        super(g, self).__init__()
        #self.k = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=1, padding=0, bias=False)

    def forward(self, z):
        a,b =torch.topk(z, 2, dim=-1, largest=True, sorted=True)
        return a
        #return self.k(a)


c = 2
h = 5
w = 5
z = torch.arange(0, c * h * w).float().view(1, c, h, w)
z.requires_grad = True
k=g()
r=k(z)
print(r)
r=r.sum()
loss=(r-1)*(r-1)
loss.backward()
print("zgrad", z.grad)

tensor([[[[ 4.,  3.],
          [ 9.,  8.],
          [14., 13.],
          [19., 18.],
          [24., 23.]],

         [[29., 28.],
          [34., 33.],
          [39., 38.],
          [44., 43.],
          [49., 48.]]]], grad_fn=<TopkBackward>)
zgrad tensor([[[[   0.,    0.,    0., 1038., 1038.],
          [   0.,    0.,    0., 1038., 1038.],
          [   0.,    0.,    0., 1038., 1038.],
          [   0.,    0.,    0., 1038., 1038.],
          [   0.,    0.,    0., 1038., 1038.]],

         [[   0.,    0.,    0., 1038., 1038.],
          [   0.,    0.,    0., 1038., 1038.],
          [   0.,    0.,    0., 1038., 1038.],
          [   0.,    0.,    0., 1038., 1038.],
          [   0.,    0.,    0., 1038., 1038.]]]])