Solved by albanD in post #2 
              
              
                Hi, 
You can use this is a toy example to see which gradient are propagated back to support with something like: 
support = torch.rand(20, requires_grad=True)
indices = # Whatever they should be

topk, indices = torch.topk(support, 10)
out = torch.zeros([B, N, N], device=self.device).scatter_(2, ind…