code: def batched_nms(
      boxes: torch.Tensor,
      scores: torch.Tensor,
      idxs: torch.Tensor,
      iou_threshold: float,
  ):
      """
      Same as torchvision.ops.boxes.batched_nms, but safer.
      """
      assert boxes.shape[-1] == 4
      # TODO may need better strategy.
      # Investigate after having a fully-cuda NMS op.
      if len(boxes) < 40000:
          return box_ops.batched_nms(boxes, scores, idxs, iou_threshold)
  
      result_mask = scores.new_zeros(scores.size(), dtype=torch.bool)
      for iid in torch.unique(idxs).cpu().tolist():
          mask = (idxs == iid).nonzero().view(-1)
          keep = nms(boxes[mask], scores[mask], iou_threshold)
          result_mask[mask[keep]] = True
      keep = result_mask.nonzero().view(-1)
      keep = keep[scores[keep].argsort(descending=True)]
      return keep
  
file_path: gluon-cv/gluoncv/torch/nn/nms.py
imports: import numpy
  import torch
  from torchvision.ops import boxes
  from torchvision.ops import nms
label_desc: Use image node_id 1 for calling a global function with example usage: batched_nms(boxes, scores, idxs, iou_threshold) and returns: keep, box_ops