// Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
#include <torch/script.h>
#include "nms_3d.h"
#include "ROIAlign_3d.h"
#include "overlap_3d.h"

using namespace at;

static auto registry =
    // torch::jit::RegisterOperators()
    torch::RegisterOperators()
    .op("nodulenet::nms_3d", &nms_3d)
    .op("nodulenet::overlap_3d", &overlap_3d)
    .op("nodulenet::roi_align_3d_forward(Tensor input, Tensor rois, float spatial_scale, int pooled_depth, int pooled_height, int pooled_width, int sampling_ratio) -> Tensor", &ROIAlign_3d_forward)
    .op("nodulenet::roi_align_3d_backward(Tensor grad, Tensor rois, float spatial_scale, int pooled_depth, int pooled_height, int pooled_width, int batch_size, int channels, int depth, int height, int width, int sampling_ratio) -> Tensor", &ROIAlign_3d_backward);
