NMS is one of those operations where it is common to write custom kernels because, as you note, implementing this in PyTorch directly is not so fast.
In the MaskRCNN-Benchmark-Implementation there is are CPU and GPU kernels for NMS, e.g.:

  
      github.com
  
  
    facebookresearch/maskrcnn-benchmark/blob/master/maskrcnn_benchmark/csrc/cpu/nms_cpu.cpp 34
// Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
#include "cpu/vision.h"


template <typename scalar_t>
at::Tensor nms_cpu_kernel(const at::Tensor& dets,
                          const at::Tensor& scores,
                          const float threshold) {
  AT_ASSERTM(!dets.type().is_cuda(), "dets must be a CPU tensor");
  AT_ASSERTM(!scores.type().is_cuda(), "scores must be a CPU tensor");
  AT_ASSERTM(dets.type() == scores.type(), "dets should have the same type as scores");

  if (dets.numel() == 0) {
    return at::empty({0}, dets.options().dtype(at::kLong).device(at::kCPU));
  }

  auto x1_t = dets.select(1, 0).contiguous();
  auto y1_t = dets.select(1, 1).contiguous();
  auto x2_t = dets.select(1, 2).contiguous();
  auto y2_t = dets.select(1, 3).contiguous();


  This file has been truncated. show original

  
  
    
    
  
  


Best regards
Thomas