add shape check when position_sensitive is true

apache · Dec 27, 2018 · 7340f14 · 7340f14
1 parent 9490fdb
commit 7340f14
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/operator/contrib/roi_align.cc b/src/operator/contrib/roi_align.cc
@@ -564,6 +564,8 @@ He, Kaiming, et al. "Mask R-CNN." ICCV, 2017
   out_shape->clear();
   if (param.position_sensitive) {
     out_shape->push_back(
+         CHECK_EQ(dshape[1] % (param.pooled_size[0]*param.pooled_size[1]), 0) <<
+           "Input channels shuold be divided by pooled_size[0]*pooled_size[1]";
          Shape4(bshape[0], dshape[1]/param.pooled_size[0]/param.pooled_size[1],
                 param.pooled_size[0], param.pooled_size[1]));
   } else {