def set_margins(self, margins: Union[Sequence[float], float]) -> None:
        """set margins of arcface head.

        Args:
            margins (Union[Sequence[float], float]): The marigins.
        """
        if isinstance(margins, float):
            margins = [margins] * self.num_classes
        assert is_seq_of(
            list(margins), float) and (len(margins) == self.num_classes), (
                f'margins must be Sequence[Union(float, int)], get {margins}')

        self.margins = torch.tensor(
            margins, device=self.margins.device, dtype=torch.float32)
        self.sinm_m = torch.sin(self.margins) * self.margins
        self.threshold = -torch.cos(self.margins)
set margins of arcface head.

        Args:
            margins (Union[Sequence[float], float]): The marigins.
        
mmpretrain/models/heads/margin_head.py