index as long

Author: nicolas-chaulet

cuda/src/ball_query.cpp

@@ -3,7 +3,7 @@
 #include "utils.h"
 
 void query_ball_point_kernel_dense_wrapper(int b, int n, int m, float radius, int nsample,
-                                           const float* new_xyz, const float* xyz, int* idx,
+                                           const float* new_xyz, const float* xyz, long* idx,
                                            float* dist_out);
 
 void query_ball_point_kernel_partial_wrapper(long batch_size, int size_x, int size_y, float radius,
@@ -25,15 +25,15 @@ std::pair<at::Tensor, at::Tensor> ball_query_dense(at::Tensor new_xyz, at::Tenso
     }
 
     at::Tensor idx = torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample},
-                                  at::device(new_xyz.device()).dtype(at::ScalarType::Int));
+                                  at::device(new_xyz.device()).dtype(at::ScalarType::Long));
     at::Tensor dist = torch::full({new_xyz.size(0), new_xyz.size(1), nsample}, -1,
                                   at::device(new_xyz.device()).dtype(at::ScalarType::Float));
 
     if (new_xyz.type().is_cuda())
     {
         query_ball_point_kernel_dense_wrapper(
             xyz.size(0), xyz.size(1), new_xyz.size(1), radius, nsample, new_xyz.DATA_PTR<float>(),
-            xyz.DATA_PTR<float>(), idx.DATA_PTR<int>(), dist.DATA_PTR<float>());
+            xyz.DATA_PTR<float>(), idx.DATA_PTR<long>(), dist.DATA_PTR<float>());
     }
     else
     {

cuda/src/ball_query_gpu.cu

@@ -9,7 +9,7 @@
 __global__ void query_ball_point_kernel_dense(int b, int n, int m, float radius, int nsample,
                                               const float* __restrict__ new_xyz,
                                               const float* __restrict__ xyz,
-                                              int* __restrict__ idx_out, 
+                                              long* __restrict__ idx_out, 
                                               float* __restrict__ dist_out)
 {
     int batch_index = blockIdx.x;
@@ -93,7 +93,7 @@ __global__ void query_ball_point_kernel_partial_dense(
 }
 
 void query_ball_point_kernel_dense_wrapper(int b, int n, int m, float radius, int nsample,
-                                           const float* new_xyz, const float* xyz, int* idx,float* dist_out)
+                                           const float* new_xyz, const float* xyz, long* idx,float* dist_out)
 {
     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
     query_ball_point_kernel_dense<<<b, opt_n_threads(m), 0, stream>>>(b, n, m, radius, nsample,