// Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
#pragma once
#include "vision.h"

at::Tensor nms_3d(const at::Tensor& dets,
                  const at::Tensor& scores,
                  const double threshold) {

    if (dets.type().is_cuda()) {
        // TODO raise error if not compiled with CUDA
        if (dets.numel() == 0)
            return at::empty({0}, dets.options().dtype(at::kLong).device(at::kCPU));
        auto b = at::cat({dets, scores.unsqueeze(1)}, 1);
        return nms_3d_cuda(b, threshold);
    }
    AT_ERROR("Not implemented on the CPU");
}
