def forward(self, predicted_heatmap, predicted_regr, target_heatmap, target_regr):
        """Compute loss for CenterNet.

        Args:
            predicted_heatmap (torch.Tensor): center heatmap prediction logits,
                expected shapes [batch, height, width, num classes].
            predicted_regr (torch.Tensor): predicted HW regression,
                expected shapes [batch, height, width, 2].
            target_heatmap ([type]): ground truth center heatmap,
                expected shapes [batch, height, width, num classes],
                each value should be in range [0,1].
            target_regr (torch.Tensor): ground truth HW regression,
                expected shapes [batch, height, width, 2].

        Returns:
            torch.Tensor with loss value.
        """
        pred_mask = torch.sigmoid(predicted_heatmap)
        mask_loss = neg_loss(pred_mask, target_heatmap)
        mask_loss *= self.mask_loss_weight

        regr_loss = (
            torch.abs(predicted_regr - target_regr).sum(1)[:, None, :, :]
            * target_heatmap
        ).sum()  # .sum(1).sum(1).sum(1)
        regr_loss = regr_loss / target_heatmap.sum()  # .sum(1).sum(1).sum(1)
        regr_loss *= self.regr_loss_weight

        loss = mask_loss + regr_loss
        if not self.size_average:
            loss *= predicted_heatmap.shape[0]

        return loss, mask_loss, regr_loss
Compute loss for CenterNet.

        Args:
            predicted_heatmap (torch.Tensor): center heatmap prediction logits,
                expected shapes [batch, height, width, num classes].
            predicted_regr (torch.Tensor): predicted HW regression,
                expected shapes [batch, height, width, 2].
            target_heatmap ([type]): ground truth center heatmap,
                expected shapes [batch, height, width, num classes],
                each value should be in range [0,1].
            target_regr (torch.Tensor): ground truth HW regression,
                expected shapes [batch, height, width, 2].

        Returns:
            torch.Tensor with loss value.
        
examples/detection/criterion.py