code: def loss_pip(
      outputs_map,
      outputs_local_x,
      outputs_local_y,
      labels_map,
      labels_local_x,
      labels_local_y,
  ):
      tmp_batch, tmp_channel, tmp_height, tmp_width = outputs_map.size()
      labels_map = labels_map.view(tmp_batch * tmp_channel, -1)
      labels_max_ids = torch.argmax(labels_map, 1)
      labels_max_ids = labels_max_ids.view(-1, 1)
  
      # print('TTT:', outputs_local_x.shape, tmp_batch, tmp_channel)
  
      outputs_local_x = outputs_local_x.reshape(
          tmp_batch * tmp_channel, -1
      )
      outputs_local_x_select = torch.gather(
          outputs_local_x, 1, labels_max_ids
      )
      outputs_local_y = outputs_local_y.reshape(
          tmp_batch * tmp_channel, -1
      )
      outputs_local_y_select = torch.gather(
          outputs_local_y, 1, labels_max_ids
      )
  
      labels_local_x = labels_local_x.view(tmp_batch * tmp_channel, -1)
      labels_local_x_select = torch.gather(
          labels_local_x, 1, labels_max_ids
      )
      labels_local_y = labels_local_y.view(tmp_batch * tmp_channel, -1)
      labels_local_y_select = torch.gather(
          labels_local_y, 1, labels_max_ids
      )
  
      labels_map = labels_map.view(
          tmp_batch, tmp_channel, tmp_height, tmp_width
      )
      loss_map = F.mse_loss(outputs_map, labels_map)
      loss_x = F.l1_loss(outputs_local_x_select, labels_local_x_select)
      loss_y = F.l1_loss(outputs_local_y_select, labels_local_y_select)
      return loss_map, loss_x, loss_y
  
file_path: insightface/reconstruction/jmlr/losses.py
imports: import torch
  from torch import nn
  import torch.nn.functional
  import kornia
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: loss_pip(outputs_map, outputs_local_x, outputs_local_y, labels_map, labels_local_x, labels_local_y) and returns: loss_map,  loss_x,  loss_y