def METHOD_NAME(self):
    batch_size = 64
    embedding_size = 32
    margin = 30
    scale = 64
    num_classes = 10
    for dtype in TEST_DTYPES:
        for sub_centers in [1, 3, 4]:
            loss_func = SubCenterArcFaceLoss(
                margin=margin,
                scale=scale,
                num_classes=num_classes,
                embedding_size=embedding_size,
                sub_centers=sub_centers,
            )
            embeddings = (
                torch.randn(batch_size, embedding_size).to(TEST_DEVICE).type(dtype)
            )
            labels = torch.randint(low=0, high=num_classes, size=(batch_size,)).to(
                TEST_DEVICE
            )
            # check if subcenters are included
            self.assertTrue(loss_func.W.shape[1] == num_classes * sub_centers)
            loss = loss_func(embeddings, labels)
            loss.backward()
            weights = F.normalize(loss_func.W, p=2, dim=0)
            logits = torch.matmul(F.normalize(embeddings), weights)
            # include only closest sub centers
            logits = logits.view(-1, num_classes, sub_centers)
            logits, _ = logits.max(axis=2)
            for i, c in enumerate(labels):
                acos = torch.acos(torch.clamp(logits[i, c], -1, 1))
                logits[i, c] = torch.cos(
                    acos
                    + torch.tensor(np.radians(margin), dtype=dtype).to(TEST_DEVICE)
                )
            correct_loss = F.cross_entropy(logits * scale, labels.to(TEST_DEVICE))
            rtol = 1e-2 if dtype == torch.float16 else 1e-5
            self.assertTrue(torch.isclose(loss, correct_loss, rtol=rtol))
            if sub_centers == 1:
                regular_arcface = ArcFaceLoss(
                    margin=margin,
                    scale=scale,
                    num_classes=num_classes,
                    embedding_size=embedding_size,
                )
                regular_arcface.W = loss_func.W
                regular_arcface_loss = regular_arcface(embeddings, labels)
                self.assertTrue(
                    torch.isclose(loss, regular_arcface_loss, rtol=rtol)
                )
            # test get_logits
            logits_out = loss_func.get_logits(embeddings)
            self.assertTrue(
                logits_out.shape == torch.Size([batch_size, num_classes])
            )
            logits = torch.matmul(F.normalize(embeddings), weights)
            # include only closest sub centers
            logits = logits.view(-1, num_classes, sub_centers)
            logits_target, _ = logits.max(axis=2)
            self.assertTrue(torch.allclose(logits_out, logits_target * scale))