Merge pull request #57 from mt1871/fix_maxpool_ceilmode

consider ceil_mode of torch.nn.MaxPool2d

Author: jaybdub

torch2trt/converters/MaxPool2d.py

@@ -1,4 +1,5 @@
 from torch2trt.torch2trt import *
+from torch2trt.module_test import add_module_test
 
 
 @tensorrt_converter('torch.nn.MaxPool2d.forward')
@@ -23,5 +24,19 @@ def convert_MaxPool2d(ctx):
         input=input._trt, type=trt.PoolingType.MAX, window_size=kernel_size)
     layer.stride = stride
     layer.padding = padding
+    if module.ceil_mode:
+        layer.padding_mode = trt.PaddingMode.EXPLICIT_ROUND_UP
 
-    output._trt = layer.get_output(0)
+    output._trt = layer.get_output(0)
+
+
+@add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 4, 6)])
+@add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 5, 7)])
+def test_MaxPool2d_without_ceil_mode():
+    return torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=False)
+
+
+@add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 4, 6)])
+@add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 5, 7)])
+def test_MaxPool2d_with_ceil_mode():
+    return torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=True)