nms.cpp 582 Bytes
Newer Older
1
#include "nms.h"
2
3
4

#include <torch/autograd.h>
#include <torch/types.h>
5
6
7
8
9
10
11
12
13
14
15
16
17
18

namespace vision {
namespace ops {

at::Tensor nms(
    const at::Tensor& dets,
    const at::Tensor& scores,
    double iou_threshold) {
  static auto op = c10::Dispatcher::singleton()
                       .findSchemaOrThrow("torchvision::nms", "")
                       .typed<decltype(nms)>();
  return op.call(dets, scores, iou_threshold);
}

19
20
21
22
TORCH_LIBRARY_FRAGMENT(torchvision, m) {
  m.def("nms(Tensor dets, Tensor scores, float iou_threshold) -> Tensor");
}

23
24
} // namespace ops
} // namespace vision