def METHOD_NAME(self, embeddings, labels):
    batch_size = labels.size(0)
    mask = torch.zeros(
        batch_size,
        self.num_classes,
        dtype=embeddings.dtype,
        device=embeddings.device,
    )
    mask[torch.arange(batch_size), labels] = 1
    return mask