code: def topk_after_batch_nms_pattern(
      cond, true_branch, data, valid_count, indices, iou_threshold
  ):
      """
      Detect the following pattern used in torchvision detection models.
  
      def batched_nms(...):
          if boxes.numel() == 0:
              return torch.empty((0,), dtype=torch.int64, device=boxes.device)
          else:
              ...
              return nms(boxes_for_nms, scores, iou_threshold)
  
      keep = batched_nms(boxes, scores, lvl, self.nms_thresh)
      keep = keep[:post_nms_top_k] # keep only topk scoring predictions
  
      An equivalent Relay subgraph:
  
      %1184 = if (%1117) {
        ...
      } else {
        ...
        %1172 = vision.non_max_suppression(%1167, %1168, %1171, -1, 0.7f, ...);
        ...
        %1183 = dyn.strided_slice(%1174, %1180, %1182, ...);
        cast(%1183, dtype="int64")
      };
      %1185 = strided_slice(%1184, begin=[0], end=[1000], strides=[1]);
  
      """
      nms = is_op("vision.non_max_suppression")(
          data, valid_count, indices, is_constant(), iou_threshold
      )
      indices = is_op("squeeze")(is_tuple_get_item(nms, 0))
      size = is_op("squeeze")(is_tuple_get_item(nms, 1))
      dyn_strided_slice = dyn_strided_slice_pattern(indices, size)
      cast_i64 = is_op("cast")(dyn_strided_slice)
  
      batched_nms_result = is_if(cond, true_branch, cast_i64)
  
      return is_op("strided_slice")(batched_nms_result)
  
imports: from ..None import expr
  from ..None import op
  from ..dataflow_pattern import wildcard, is_constant, is_op, rewrite, is_tuple, is_tuple_get_item, is_if, DFPatternCallback
label_desc: Use image node_id 5 for calling a global function with example usage: topk_after_batch_nms_pattern(cond, true_branch, data, valid_count, indices, iou_threshold) and returns: is_op