TaoLv commented on a change in pull request #17884: [MKL-DNN] Integrate Conv3d 
and Pool3d/1d
URL: https://github.com/apache/incubator-mxnet/pull/17884#discussion_r397866258
 
 

 ##########
 File path: src/operator/nn/mkldnn/mkldnn_pooling-inl.h
 ##########
 @@ -114,15 +115,21 @@ inline bool SupportMKLDNNPooling(const PoolingParam 
&param,
     return true;
   } else {
     if (param.pool_type == pool_enum::kAvgPooling) {
-      CHECK_EQ(dshape.ndim(), 4);
+      CHECK(dshape.ndim() == 3 || dshape.ndim() == 4 || dshape.ndim() == 5);
       // mkldnn works differently when padding is asymmetric, so let's skip 
this case.
-      if (param.pad[0] == GetPaddingSizeFull(dshape[2], param.pad[0], 
param.pad[0], param.kernel[0],
-                                             param.stride[0]) &&
-          param.pad[1] == GetPaddingSizeFull(dshape[3], param.pad[1], 
param.pad[1], param.kernel[1],
-                                             param.stride[1])) {
-        return true;
+      bool is_symmetric = true;
+      switch (dshape.ndim()) {
+        case 5:
+          is_symmetric = is_symmetric && (param.pad[2] == 
GetPaddingSizeFull(dshape[4],
+                                param.pad[2], param.pad[2], param.kernel[2], 
param.stride[2]));
+        case 4:
+          is_symmetric = is_symmetric && (param.pad[1] == 
GetPaddingSizeFull(dshape[3],
+                                param.pad[1], param.pad[1], param.kernel[1], 
param.stride[1]));
 
 Review comment:
   okay, i didn't realize that you don't have `break` for each case. So if 
`ndim == 4`, both pad[1] and pad[0] are checked.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


With regards,
Apache Git Services

Reply via email to