code: def forward(self, output, records):
      target = records["target"]
      target_weight = records["target_weight"]
      batch_size = output.shape[0]
      num_joints = output.shape[1]
      heatmaps_pred = output.reshape(
          (batch_size, num_joints, -1)
      ).split(num_joints, 1)
      heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(
          num_joints, 1
      )
      loss = 0
      for idx in range(num_joints):
          heatmap_pred = heatmaps_pred[idx].squeeze()
          heatmap_gt = heatmaps_gt[idx].squeeze()
          if self.use_target_weight:
              loss += self.loss_scale * self.criterion(
                  heatmap_pred.multiply(target_weight[:, idx]),
                  heatmap_gt.multiply(target_weight[:, idx]),
              )
          else:
              loss += self.loss_scale * self.criterion(
                  heatmap_pred, heatmap_gt
              )
      keypoint_losses = dict()
      keypoint_losses["loss"] = loss / num_joints
      return keypoint_losses
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from itertools import cycle, islice
  from collections import abc
  import numpy
  import paddle
  import paddle.nn
  from ppdet.core.workspace import register, serializable
label_desc: Use image node_id 2 for calling the KeyPointMSELoss obj's underlying member method code with example usage: obj.forward(output, records) and returns: keypoint_losses