Move voxel_size to the device where points are (open-mmlab#2681)

* ... * ... * ...
CokeDong · Apr 6, 2023 · 665d99d · 665d99d
1 parent 88a1f9b
commit 665d99d
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/mmcv/ops/csrc/pytorch/npu/voxelization_npu.cpp b/mmcv/ops/csrc/pytorch/npu/voxelization_npu.cpp
@@ -24,13 +24,11 @@ int hard_voxelize_forward_npu(const at::Tensor &points, at::Tensor &voxels,
 
   at::Tensor voxel_size_cpu = at::from_blob(
       const_cast<float *>(voxel_size.data()), {3}, dtype(at::kFloat));
-  at::Tensor voxel_size_npu =
-      CalcuOpUtil::CopyTensorHostToDevice(voxel_size_cpu);
+  at::Tensor voxel_size_npu = voxel_size_cpu.to(points.device());
 
   at::Tensor coors_range_cpu = at::from_blob(
       const_cast<float *>(coors_range.data()), {6}, dtype(at::kFloat));
-  at::Tensor coors_range_npu =
-      CalcuOpUtil::CopyTensorHostToDevice(coors_range_cpu);
+  at::Tensor coors_range_npu = coors_range_cpu.to(points.device());
 
   int64_t max_points_ = (int64_t)max_points;
   int64_t max_voxels_ = (int64_t)max_voxels;