From b55dee9fe15d679f9d92ec660d77059ff380c839 Mon Sep 17 00:00:00 2001 From: lixinyu Date: Tue, 7 Apr 2020 01:09:23 -0700 Subject: [PATCH] fix max_pool2d cuda version Dimension out of range issue(#36046) (#36095) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/36095 Test Plan: Imported from OSS Differential Revision: D20876733 Pulled By: glaringlee fbshipit-source-id: a2b92fd2dd0254c5443af469e3fb2faa2323e5c9 --- aten/src/ATen/native/cuda/DilatedMaxPool2d.cu | 4 ++-- torch/testing/_internal/common_nn.py | 10 +++++++++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/aten/src/ATen/native/cuda/DilatedMaxPool2d.cu b/aten/src/ATen/native/cuda/DilatedMaxPool2d.cu index 57b1f5371b67..30c51b8db025 100644 --- a/aten/src/ATen/native/cuda/DilatedMaxPool2d.cu +++ b/aten/src/ATen/native/cuda/DilatedMaxPool2d.cu @@ -358,7 +358,7 @@ void max_pool2d_with_indices_out_cuda_template( Tensor input = input_.contiguous(memory_format); - const int64_t in_stride_n = input.stride(-4); + const int64_t in_stride_n = input_.ndimension() == 4 ? input.stride(-4) : 0; const int64_t in_stride_c = input.stride(-3); const int64_t in_stride_h = input.stride(-2); const int64_t in_stride_w = input.stride(-1); @@ -506,7 +506,7 @@ void max_pool2d_with_indices_backward_out_cuda_template( const int64_t inputHeight = input.size(-2); const int64_t inputWidth = input.size(-1); - const int64_t in_stride_n = input.stride(-4); + const int64_t in_stride_n = input.ndimension() == 4 ? input.stride(-4) : 0; const int64_t in_stride_c = input.stride(-3); const int64_t in_stride_h = input.stride(-2); const int64_t in_stride_w = input.stride(-1); diff --git a/torch/testing/_internal/common_nn.py b/torch/testing/_internal/common_nn.py index 4e8e022d91bf..3daf395825c4 100644 --- a/torch/testing/_internal/common_nn.py +++ b/torch/testing/_internal/common_nn.py @@ -66,7 +66,7 @@ def get_weight(m): # and the `cpp_var_map` entry must be # `{'random_samples': random_samples}` in order to populate the C++ variable `random_samples` # used in the C++ constructor argument with the Python tensor value `random_samples`. -# +# # For NN functional: # 1. Make sure you already have a test dict with the functional configuration you want to test. # 2. If the test dict's `constructor` entry looks like `wrap_functional(F.some_functional_name, ...)`, @@ -1812,12 +1812,20 @@ def fractional_max_pool3d_test(test_case): cpp_constructor_args='torch::nn::Conv2dOptions(4, 4, {2, 2}).dilation({2, 2}).groups(4)', input_size=(2, 4, 5, 5), ), + dict( + module_name='MaxPool2d', + constructor_args=((3, 3), (2, 2), (1, 1)), + cpp_constructor_args='torch::nn::MaxPool2dOptions({3, 3}).stride({2, 2}).padding({1, 1})', + input_size=(3, 7, 7), + desc='3d_input' + ), dict( module_name='MaxPool2d', constructor_args=((3, 3), (2, 2), (1, 1)), cpp_constructor_args='torch::nn::MaxPool2dOptions({3, 3}).stride({2, 2}).padding({1, 1})', input_size=(1, 3, 7, 7), check_with_channels_last=True, + desc='4d_input' ), dict( module_name='AvgPool1d',