def _ae_loss_per_image(self, tags: Tensor, keypoint_indices: Tensor):
        """Compute associative embedding loss for one image.

        Args:
            tags (Tensor): Tagging heatmaps in shape (K*L, H, W)
            keypoint_indices (Tensor): Ground-truth keypint position indices
                in shape (N, K, 2)
        """
        K = keypoint_indices.shape[1]
        C, H, W = tags.shape
        L = C // K

        tags = tags.view(L, K, H * W)
        instance_tags = []
        instance_kpt_tags = []

        for keypoint_indices_n in keypoint_indices:
            _kpt_tags = []
            for k in range(K):
                if keypoint_indices_n[k, 1]:
                    _kpt_tags.append(tags[:, k, keypoint_indices_n[k, 0]])

            if _kpt_tags:
                kpt_tags = torch.stack(_kpt_tags)
                instance_kpt_tags.append(kpt_tags)
                instance_tags.append(kpt_tags.mean(dim=0))

        N = len(instance_kpt_tags)  # number of instances with valid keypoints

        if N == 0:
            pull_loss = tags.new_zeros(size=(), requires_grad=True)
            push_loss = tags.new_zeros(size=(), requires_grad=True)
        else:
            pull_loss = sum(
                F.mse_loss(_kpt_tags, _tag.expand_as(_kpt_tags))
                for (_kpt_tags, _tag) in zip(instance_kpt_tags, instance_tags))

            if N == 1:
                push_loss = tags.new_zeros(size=(), requires_grad=True)
            else:
                tag_mat = torch.stack(instance_tags)  # (N, L)
                diff = tag_mat[None] - tag_mat[:, None]  # (N, N, L)
                push_loss = torch.sum(torch.exp(-diff.pow(2)))

            # normalization
            eps = 1e-6
            pull_loss = pull_loss / (N + eps)
            push_loss = push_loss / ((N - 1) * N + eps)

        return pull_loss, push_loss
Compute associative embedding loss for one image.

        Args:
            tags (Tensor): Tagging heatmaps in shape (K*L, H, W)
            keypoint_indices (Tensor): Ground-truth keypint position indices
                in shape (N, K, 2)
        
mmpose/mmpose/models/losses/ae_loss.py