From 7340f14f84223c3487893a72b1f26c930d56dec8 Mon Sep 17 00:00:00 2001
From: shesung <shesung.sam@gmail.com>
Date: Thu, 27 Dec 2018 10:25:41 +0800
Subject: [PATCH] add shape check when position_sensitive is true

---
 src/operator/contrib/roi_align.cc | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/operator/contrib/roi_align.cc b/src/operator/contrib/roi_align.cc
index 3333c9539307..b32b5f0fee66 100644
--- 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 {