fix roi_pool bug.

open-mmlab · Jun 11, 2024 · 65e40cf · 65e40cf
1 parent ba70ee3
commit 65e40cf
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/mmcv/ops/csrc/pytorch/npu/roi_pool_npu.cpp b/mmcv/ops/csrc/pytorch/npu/roi_pool_npu.cpp
@@ -70,7 +70,8 @@ void roi_pool_backward_npu(Tensor grad_output, Tensor rois, Tensor argmax,
       .Attr("spatial_scale_w", spatial_scale)
       .Attr("pool_channel", pooled_channel)
       .Run();
-  at::Tensor res = y.contiguous();
+  at::Tensor result = y.transpose(2, 3).transpose(1, 2);
+  at::Tensor res = result.contiguous();
   grad_input.copy_(res);
 }