radius_cpu.h 448 Bytes
Newer Older
Alexander Liao's avatar
Alexander Liao committed
1
2
3
4
5
6
7
#pragma once

#include <torch/extension.h>
#include "utils/neighbors.cpp"
#include <iostream>
#include "compat.h"

8
torch::Tensor radius_cpu(torch::Tensor query, torch::Tensor support,
9
			 			 double radius, int64_t max_num, int64_t n_threads);
10
11
12
13
14

torch::Tensor batch_radius_cpu(torch::Tensor query,
			       torch::Tensor support,
			       torch::Tensor query_batch,
			       torch::Tensor support_batch,
15
			       double radius, int64_t max_num);