def forward(self, tags: Tensor, keypoint_indices: Union[List[Tensor],
                                                            Tensor]):
        """Compute associative embedding loss on a batch of data.

        Args:
            tags (Tensor): Tagging heatmaps in shape (B, L*K, H, W)
            keypoint_indices (Tensor|List[Tensor]): Ground-truth keypint
                position indices represented by a Tensor in shape
                (B, N, K, 2), or a list of B Tensors in shape (N_i, K, 2)
                Each keypoint's index is represented as [i, v], where i is the
                position index in the heatmap (:math:`i=y*w+x`) and v is the
                visibility

        Returns:
            tuple:
            - pull_loss (Tensor)
            - push_loss (Tensor)
        """

        assert tags.shape[0] == len(keypoint_indices)

        pull_loss = 0.
        push_loss = 0.

        for i in range(tags.shape[0]):
            _pull, _push = self._ae_loss_per_image(tags[i],
                                                   keypoint_indices[i])
            pull_loss += _pull * self.loss_weight
            push_loss += _push * self.loss_weight * self.push_loss_factor

        return pull_loss, push_loss
Compute associative embedding loss on a batch of data.

        Args:
            tags (Tensor): Tagging heatmaps in shape (B, L*K, H, W)
            keypoint_indices (Tensor|List[Tensor]): Ground-truth keypint
                position indices represented by a Tensor in shape
                (B, N, K, 2), or a list of B Tensors in shape (N_i, K, 2)
                Each keypoint's index is represented as [i, v], where i is the
                position index in the heatmap (:math:`i=y*w+x`) and v is the
                visibility

        Returns:
            tuple:
            - pull_loss (Tensor)
            - push_loss (Tensor)
        
mmpose/mmpose/models/losses/ae_loss.py