Fix a bug in pooling.

cjolivier01 · Feb 2, 2018 · 4742ad5 · 4742ad5
1 parent 6222162
commit 4742ad5
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/operator/nn/pooling.cc b/src/operator/nn/pooling.cc
@@ -279,8 +279,9 @@ inline static bool PoolingStorageType(const nnvm::NodeAttrs &attrs,
     return storage_type_assign(out_attrs, mxnet::kDefaultStorage,
                                dispatch_mode, DispatchMode::kFComputeEx);
   }
-#endif
+#else
   CHECK_EQ(out_attrs->size(), 1);
+#endif
   return storage_type_assign(out_attrs, mxnet::kDefaultStorage,
                              dispatch_mode, DispatchMode::kFCompute);
 }
@@ -299,8 +300,9 @@ inline static bool BackwardPoolingStorageType(const nnvm::NodeAttrs &attrs,
     return storage_type_assign(out_attrs, mxnet::kDefaultStorage,
                                dispatch_mode, DispatchMode::kFComputeEx);
   }
-#endif
+#else
   CHECK_EQ(in_attrs->size(), 3);
+#endif
   return storage_type_assign(out_attrs, mxnet::kDefaultStorage,
                              dispatch_mode, DispatchMode::kFCompute);
 }