def _get_logit_with_margin(self, pre_logits, target):
        """add arc margin to the cosine in target index.

        The target must be in index format.
        """
        assert target.dim() == 1 or (
            target.dim() == 2 and target.shape[1] == 1), \
            'The target must be in index format.'
        cosine = self.norm_product(pre_logits)
        phi = torch.cos(torch.acos(cosine) + self.margins)

        if self.easy_margin:
            # when cosine>0, choose phi
            # when cosine<=0, choose cosine
            phi = torch.where(cosine > 0, phi, cosine)
        else:
            # when cos>th, choose phi
            # when cos<=th, choose cosine-mm
            phi = torch.where(cosine > self.threshold, phi,
                              cosine - self.sinm_m)

        target = convert_to_one_hot(target, self.num_classes)
        output = target * phi + (1 - target) * cosine
        return output
add arc margin to the cosine in target index.

        The target must be in index format.
        
mmpretrain/models/heads/margin_head.py