code: def batched_nms(boxes, scores, idxs, iou_threshold):
      # type: (Tensor, Tensor, Tensor, float) -> Tensor
      """
      Performs non-maximum suppression in a batched fashion.
  
      Each index value correspond to a category, and NMS
      will not be applied between elements of different categories.
  
      Parameters
      ----------
      boxes : Tensor[N, 4]
          boxes where NMS will be performed. They
          are expected to be in (x1, y1, x2, y2) format
      scores : Tensor[N]
          scores for each one of the boxes
      idxs : Tensor[N]
          indices of the categories for each one of the boxes.
      iou_threshold : float
          discards all overlapping boxes
          with IoU < iou_threshold
  
      Returns
      -------
      keep : Tensor
          int64 tensor with the indices of
          the elements that have been kept by NMS, sorted
          in decreasing order of scores
      """
      if boxes.numel() == 0:
          return torch.empty(
              (0,), dtype=torch.int64, device=boxes.device
          )
  
      # strategy: in order to perform NMS independently per class.
      # we add an offset to all the boxes. The offset is dependent
      # only on the class idx, and is large enough so that boxes
      # from different classes do not overlap
      # 获取所有boxes中最大的坐标值（xmin, ymin, xmax, ymax）
      max_coordinate = boxes.max()
  
      # to(): Performs Tensor dtype and/or device conversion
      # 为每一个类别生成一个很大的偏移量
      # 这里的to只是让生成tensor的dytpe和device与boxes保持一致
      offsets = idxs.to(boxes) * (max_coordinate + 1)
      # boxes加上对应层的偏移量后，保证不同类别之间boxes不会有重合的现象
      boxes_for_nms = boxes + offsets[:, None]
      keep = nms(boxes_for_nms, scores, iou_threshold)
      return keep
  
imports: from math import sqrt
  import itertools
  import torch
  import torch.nn.functional
  from torch.jit.annotations import Tuple, List
  from torch import nn, Tensor
  import numpy
label_desc: Use image node_id 5 for calling a global function with example usage: batched_nms(boxes, scores, idxs, iou_threshold) and returns: keep, torch