Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,089 Bytes
938e515 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
// Copyright (c) Facebook, Inc. and its affiliates.
#pragma once
#include <torch/types.h>
namespace detectron2 {
at::Tensor nms_rotated_cpu(
const at::Tensor& dets,
const at::Tensor& scores,
const double iou_threshold);
#if defined(WITH_CUDA) || defined(WITH_HIP)
at::Tensor nms_rotated_cuda(
const at::Tensor& dets,
const at::Tensor& scores,
const double iou_threshold);
#endif
// Interface for Python
// inline is needed to prevent multiple function definitions when this header is
// included by different cpps
inline at::Tensor nms_rotated(
const at::Tensor& dets,
const at::Tensor& scores,
const double iou_threshold) {
assert(dets.device().is_cuda() == scores.device().is_cuda());
if (dets.device().is_cuda()) {
#if defined(WITH_CUDA) || defined(WITH_HIP)
return nms_rotated_cuda(
dets.contiguous(), scores.contiguous(), iou_threshold);
#else
AT_ERROR("Detectron2 is not compiled with GPU support!");
#endif
}
return nms_rotated_cpu(dets.contiguous(), scores.contiguous(), iou_threshold);
}
} // namespace detectron2
|