29 #include "torch/script.h" 38 : device_type(device_type), device_idx(device_idx) {}
41 neighbors_index = torch::empty(
43 .device(device_type, device_idx));
44 *ptr = neighbors_index.data_ptr<
int32_t>();
48 neighbors_distance = torch::empty(
49 {int64_t(num)}, torch::dtype(ToTorchDtype<T>())
50 .device(device_type, device_idx));
51 *ptr = neighbors_distance.data_ptr<T>();
55 return neighbors_index.data_ptr<
int32_t>();
58 const T*
DistancesPtr()
const {
return neighbors_distance.data_ptr<T>(); }
62 return neighbors_distance;
66 torch::Tensor neighbors_index;
67 torch::Tensor neighbors_distance;
68 torch::DeviceType device_type;
const torch::Tensor & NeighborsIndex() const
Definition: NeighborSearchAllocator.h:60
const int32_t * IndicesPtr() const
Definition: NeighborSearchAllocator.h:54
const char const char value recording_handle imu_sample recording_handle uint8_t size_t data_size k4a_record_configuration_t config target_format k4a_capture_t capture_handle k4a_imu_sample_t imu_sample playback_handle k4a_logging_message_cb_t void min_level device_handle k4a_imu_sample_t int32_t
Definition: K4aPlugin.cpp:398
void AllocDistances(T **ptr, size_t num)
Definition: NeighborSearchAllocator.h:47
NeighborSearchAllocator(torch::DeviceType device_type, int device_idx)
Definition: NeighborSearchAllocator.h:37
void AllocIndices(int32_t **ptr, size_t num)
Definition: NeighborSearchAllocator.h:40
Definition: NeighborSearchAllocator.h:35
const T * DistancesPtr() const
Definition: NeighborSearchAllocator.h:58
TorchDtype_t ToTorchDtype< int32_t >()
Definition: TorchHelper.h:90
const torch::Tensor & NeighborsDistance() const
Definition: NeighborSearchAllocator.h:61