code: def batched_nms(
      boxes: Tensor,
      scores: Tensor,
      idxs: Tensor,
      iou_threshold: 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.
  
      Args:
          boxes (Tensor[N, 4]): boxes where NMS will be performed. They
              are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and
              ``0 <= y1 < y2``.
          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:
          Tensor: int64 tensor with the indices of the elements that have been kept by NMS, sorted
          in decreasing order of scores
      """
      if not torch.jit.is_scripting() and not torch.jit.is_tracing():
          _log_api_usage_once(batched_nms)
      # Benchmarks that drove the following thresholds are at
      # https://github.com/pytorch/vision/issues/1311#issuecomment-781329339
      if (
          boxes.numel()
          > (4000 if boxes.device.type == "cpu" else 20000)
          and not torchvision._is_tracing()
      ):
          return _batched_nms_vanilla(
              boxes, scores, idxs, iou_threshold
          )
      else:
          return _batched_nms_coordinate_trick(
              boxes, scores, idxs, iou_threshold
          )
  
imports: from typing import Tuple
  import torch
  import torchvision
  from torch import Tensor
  from torchvision.extension import _assert_has_ops
  from ..utils import _log_api_usage_once
  from ._box_convert import _box_cxcywh_to_xyxy, _box_xywh_to_xyxy, _box_xyxy_to_cxcywh, _box_xyxy_to_xywh
  from ._utils import _upcast
label_desc: Use image node_id 2 for calling a global function with example usage: batched_nms(boxes, scores, idxs, iou_threshold) and returns: _batched_nms_vanilla, _batched_nms_coordinate_trick