[Fix] Fix ball_query precision bug fp16 to fp32 (#3110)

open-mmlab · May 22, 2024 · c0c81ce · c0c81ce
1 parent 27df2b3
commit c0c81ce
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mmcv/ops/csrc/pytorch/npu/ball_query_npu.cpp b/mmcv/ops/csrc/pytorch/npu/ball_query_npu.cpp
@@ -9,10 +9,10 @@ void ball_query_forward_npu(int b, int n, int m, float min_radius,
   int64_t nsample_i64 = nsample;
 
   // transpose new_xyz from [B, M, 3] to [M, B, 3]
-  at::Tensor new_xyz_transpose = new_xyz.transpose(0, 1);
+  at::Tensor new_xyz_transpose = new_xyz.transpose(0, 1).to(at::kFloat);
 
   // transpose xyz from [B, N, 3] to [B, 3, N]
-  at::Tensor xyz_transpose = xyz.transpose(1, 2);
+  at::Tensor xyz_transpose = xyz.transpose(1, 2).to(at::kFloat);
 
   // transpose idx from [B, M, nsample] to [M, B, nsample]
   at::Tensor idx_transpose = idx.transpose(0, 1).contiguous();