How can I get gradients of ‘Li’ w.r.t. its corresponding ‘Zi_unnormalized’ and such that a tensor of shape same as that of feat is formed with these gradients?
    ....
    @staticmethod
    def backward(ctx, grad_output):    # grad_output size = ([batch_size, 128,4,4])
        
        input, weight, bias, output = ctx.saved_tensors    #input size = ([batch_size, 96,8,8])
      
        feat = output.clone().detach()   # output from forward with size = ([batch_size, 128,4,4])

        features = feat.view(feat.shape[0], feat.shape[1], -1) # features size = ([batch_size, 128,16])

        i = 0        
        cont = torch.tensor([0.]).to(dev)
        while i in range(0, features.shape[0]):
              for f in range(len(features[i])):
                  Zi_unnormalized = features[i][f].requires_grad_(requires_grad=True)
                  Zi = torch.nn.functional.normalize(Zi_unnormalized, dim = 0)
                  # Zj and Zk are tensors made from features[i][*] and features[other than i][*]. Zj and Zk varies for each Zi (or f)

                  Zi_Zk = torch.Tensor([0]).to(dev)
                  for k in Zk:
                      k= torch.nn.functional.normalize(k, dim = 0)
                      zi_zk = torch.exp(torch.div(torch.dot(Zi,k.T), 0.07))
                      Zi_Zk = Zi_Zk.add(zi_zk)

                  # Similarly computing Zi_Zj
                  # Li = some algebra of Zi_Zj and Zi_Zk
                  # number of 'Li' values =  features.shape[0] * features.shape[1]
                  cont = cont.add(Li)   # 1 value
              i+=1

I tried autograd backward but I get every time a RuntimeError. Then I saw that the grad_fn of Zi everytime remains None!!