text: import torch.nn as nn
  import torch.nn.functional as F
  
  from rfvision.models.builder import LOSSES
  from .utils import weighted_loss
  
  @LOSSES.register_module()
  class JointsMSELoss(nn.Module):
      """MSE loss for heatmaps.
  
      Args:
          use_target_weight (bool): Option to use weighted MSE loss.
              Different joint types may have different target weights.
      """
  
      def __init__(self, use_target_weight=False):
          super().__init__()
          self.criterion = nn.MSELoss()
          self.use_target_weight = use_target_weight
  
      def forward(self, output, target, target_weight):
          """Forward function."""
          batch_size = output.size(0)
          num_joints = output.size(1)
  
          heatmaps_pred = output.reshape(
              (batch_size, num_joints, -1)).split(1, 1)
          heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)
  
          loss = 0.
  
          for idx in range(num_joints):
              heatmap_pred = heatmaps_pred[idx].squeeze(1)
              heatmap_gt = heatmaps_gt[idx].squeeze(1)
              if self.use_target_weight:
                  loss += self.criterion(
                      heatmap_pred.mul(target_weight[:, idx]),
                      heatmap_gt.mul(target_weight[:, idx]))
              else:
                  loss += self.criterion(heatmap_pred, heatmap_gt)
  
          return loss / num_joints
  
  @weighted_loss
  def mse_loss(pred, target):
      """Warpper of mse loss."""
      return F.mse_loss(pred, target, reduction='none')
  
  
  @LOSSES.register_module()
  class MSELoss(nn.Module):
      """MSELoss.
  
      Args:
          reduction (str, optional): The method that reduces the loss to a
              scalar. Options are "none", "mean" and "sum".
          loss_weight (float, optional): The weight of the loss. Defaults to 1.0
      """
  
      def __init__(self, reduction='mean', loss_weight=1.0):
          super().__init__()
          self.reduction = reduction
          self.loss_weight = loss_weight
  
      def forward(self, pred, target, weight=None, avg_factor=None):
          """Forward function of loss.
  
          Args:
              pred (torch.Tensor): The prediction.
              target (torch.Tensor): The learning target of the prediction.
              weight (torch.Tensor, optional): Weight of the loss for each
                  prediction. Defaults to None.
              avg_factor (int, optional): Average factor that is used to average
                  the loss. Defaults to None.
  
          Returns:
              torch.Tensor: The calculated loss
          """
          loss = self.loss_weight * mse_loss(
              pred,
              target,
              weight,
              reduction=self.reduction,
              avg_factor=avg_factor)
          return loss
  
  