solve deformable_col2im: invalid argument by increasing channel

xvjiarui · Apr 3, 2019 · 7307e40 · 7307e40
1 parent edb0393
commit 7307e40
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/mmdet/ops/dcn/src/deform_conv_cuda_kernel.cu b/mmdet/ops/dcn/src/deform_conv_cuda_kernel.cu
@@ -73,10 +73,11 @@ using namespace at;
        i += blockDim.x * gridDim.x)
 
 const int CUDA_NUM_THREADS = 1024;
+const int kMaxGridNum = 65535;
 
 inline int GET_BLOCKS(const int N)
 {
-  return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;
+  return std::min(kMaxGridNum, (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS);
 }
 
 template <typename scalar_t>