fix(aten::max_pool2d): Supressing error due to not filling in stride in

the default case Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
pytorch · Nov 19, 2020 · ed3c185 · ed3c185
1 parent 4ee6c20
commit ed3c185
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/core/conversion/converters/impl/pooling.cpp b/core/conversion/converters/impl/pooling.cpp
@@ -30,6 +30,10 @@ bool MaxPoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, args& ar
   auto padding = util::toDims(args[3].unwrapToIntList());
   LOG_DEBUG("padding: " << padding);
   auto stride = util::toDims(args[2].unwrapToIntList());
+  if (args[2].unwrapToIntList().size() == 0) {
+    LOG_DEBUG("Stride not providied, using kernel_size as stride");
+    stride = util::toDims(args[1].unwrapToIntList());
+  }
   LOG_DEBUG("stride: " << stride);
 
   auto dilation = util::toDims(args[4].unwrapToIntList());
@@ -88,6 +92,10 @@ bool AvgPoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, args& ar
   auto padding = util::toDims(args[3].unwrapToIntList());
   LOG_DEBUG("padding: " << padding);
   auto stride = util::toDims(args[2].unwrapToIntList());
+  if (args[2].unwrapToIntList().size() == 0) {
+    LOG_DEBUG("Stride not providied, using kernel_size as stride");
+    stride = util::toDims(args[1].unwrapToIntList());
+  }
   LOG_DEBUG("stride: " << stride);
 
   bool ceil_mode = args[4].unwrapToBool();