更新knn和three_nn的NPU适配代码 (#3194)

open-mmlab · Nov 18, 2024 · e1aab12 · e1aab12
1 parent 71437a3
commit e1aab12
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/mmcv/ops/csrc/pytorch/npu/knn_npu.cpp b/mmcv/ops/csrc/pytorch/npu/knn_npu.cpp
@@ -12,7 +12,7 @@ void knn_forward_npu(int b, int n, int m, int nsample, const Tensor xyz,
   at::Tensor target = new_xyz.contiguous();
 
   bool is_from_knn = true;
-  EXEC_NPU_CMD(aclnnKnn, source, target, nsample, is_from_knn, idx, dist2);
+  EXEC_NPU_CMD(aclnnKnn, source, target, is_from_knn, nsample, dist2, idx);
 }
 
 void knn_forward_impl(int b, int n, int m, int nsample, const Tensor xyz,

diff --git a/mmcv/ops/csrc/pytorch/npu/three_nn_npu.cpp b/mmcv/ops/csrc/pytorch/npu/three_nn_npu.cpp
@@ -18,7 +18,7 @@ void three_nn_forward_npu(int b, int n, int m, const Tensor unknown,
 
   bool is_from_knn = false;
   uint32_t nsample = 3;
-  EXEC_NPU_CMD(aclnnKnn, source, target, nsample, is_from_knn, idx, dist2);
+  EXEC_NPU_CMD(aclnnKnn, source, target, is_from_knn, nsample, dist2, idx);
   if (originDtype == at::kHalf) {
     dist2 = dist2.to(at::kHalf);
   }