class ArcMarginProduct(nn.Module):
    r"""Implement of large margin arc distance: :
        Args:
            in_features: size of each input sample
            out_features: size of each output sample
            s: norm of input feature
            m: margin
            cos(theta + m)
        """
    def __init__(self, in_features, out_features, s=30.0, m=0.50, easy_margin=False, ls_eps=0.0):
        super(ArcMarginProduct, self).__init__()
        self.in_features = torch.tensor(in_features)
        self.out_features = torch.tensor(out_features)
        self.s = torch.tensor(s)
        self.m = torch.repeat_interleave(torch.tensor([m],requires_grad=True),CFG.num_classes)
        self.ls_eps = torch.tensor(ls_eps)  # label smoothing
        self.weight = Parameter(torch.FloatTensor(out_features, in_features))
        nn.init.xavier_uniform_(self.weight)
        self.easy_margin = easy_margin
        self.torch_pi = torch.acos(torch.zeros(1))*2
#         self.critition = critition
    def forward(self, input, label, critition=None):
        
        m = self.m[label]
        cos_m = torch.tensor([torch.cos(value) for value in m]).to(CFG.device).reshape(-1,1)
        sin_m = torch.tensor([torch.sin(value) for value in m]).to(CFG.device).reshape(-1,1)
        th = torch.tensor([torch.cos(torch.sub(self.torch_pi, value)) for value in m]).to(CFG.device).reshape(-1,1)
        mm = torch.tensor([torch.mul(torch.sin(torch.sub(self.torch_pi, value)), value) for value in m]).to(CFG.device).reshape(-1,1)
        
        # --------------------------- cos(theta) & phi(theta) ---------------------------
        cosine = F.linear(F.normalize(input), F.normalize(self.weight))
        sine = torch.sqrt(torch.sub(1.0,torch.pow(cosine, 2)))
        phi = torch.sub(torch.mul(cosine,cos_m),torch.mul(sine,sin_m))
        
        if self.easy_margin:
            phi = torch.where(cosine > 0, phi, cosine)
        else:
            phi = torch.where(cosine > th, phi, torch.sub(cosine,mm))
        # --------------------------- convert label to one-hot ---------------------------
        # one_hot = torch.zeros(cosine.size(), requires_grad=True, device='cuda')
        one_hot = torch.zeros(cosine.size(), device='cuda')
        one_hot.scatter_(1, label.view(-1, 1).long(), 1)
        if self.ls_eps > 0:
            one_hot = torch.add(torch.mul(torch.sub(1, self.ls_eps), one_hot), torch.div(self.ls_eps, self.out_features))
        # -------------torch.where(out_i = {x_i if condition_i else y_i) -------------
        output = torch.add(torch.mul(one_hot, phi), torch.mul(torch.sub(1.0,one_hot),cosine))
        output = torch.mul(output, self.s)
        
        ### update m
        if critition != None:
#             training_gradient = (training_loss)
            loss = critition(output, label)
            out = torch.autograd.grad(loss, m , create_graph=True, retain_graph=True,allow_unused=True)
            print(out) # output none
#             self.m[label] = self.m[label] - phi_b[0]
            return output, loss
        return output

I replace all operator *,-,/,+ to torch pointwise, and assing m = self.m[index] and resuse it. And I got the same error, output of torch grad equal “None”.  . Thank you for your work.