nms.h 251 Bytes
Newer Older
1
#pragma once
Vasilis Vryniotis's avatar
Vasilis Vryniotis committed
2

3
#include <ATen/ATen.h>
4
#include "../macros.h"
5

6
7
8
namespace vision {
namespace ops {

9
VISION_API at::Tensor nms(
10
11
    const at::Tensor& dets,
    const at::Tensor& scores,
12
13
14
15
    double iou_threshold);

} // namespace ops
} // namespace vision