code: def get_max_pred(batch_heatmaps):
      batch_size = batch_heatmaps.shape[0]
      num_joints = batch_heatmaps.shape[1]
      width = batch_heatmaps.shape[3]
      heatmaps_reshaped = batch_heatmaps.reshape(
          (batch_size, num_joints, -1)
      )
      idx = nd.argmax(heatmaps_reshaped, 2)
      maxvals = nd.max(heatmaps_reshaped, 2)
  
      maxvals = maxvals.reshape((batch_size, num_joints, 1))
      idx = idx.reshape((batch_size, num_joints, 1))
  
      preds = nd.tile(idx, (1, 1, 2)).astype(np.float32)
  
      preds[:, :, 0] = (preds[:, :, 0]) % width
      preds[:, :, 1] = nd.floor((preds[:, :, 1]) / width)
  
      pred_mask = nd.tile(nd.greater(maxvals, 0.0), (1, 1, 2))
      pred_mask = pred_mask.astype(np.float32)
  
      preds *= pred_mask
      return preds, maxvals
  
imports: from __future__ import absolute_import
  from __future__ import division
  import math
  import random
  import numpy
  import mxnet
  from mxnet import nd, image
  from mxnet.gluon.data.vision import transforms
  from ...utils.filesystem import try_import_cv2
label_desc: Use image node_id 10 for calling a global function with example usage: get_max_pred(batch_heatmaps) and returns: preds,  maxvals