diff --git a/src/frontends/pytorch/src/op/adaptive_max_pool2d.cpp b/src/frontends/pytorch/src/op/adaptive_max_pool2d.cpp index 77e6b41ecef4de..134434330cca28 100644 --- a/src/frontends/pytorch/src/op/adaptive_max_pool2d.cpp +++ b/src/frontends/pytorch/src/op/adaptive_max_pool2d.cpp @@ -14,13 +14,8 @@ namespace op { OutputVector translate_adaptive_max_pool2d(NodeContext& context) { auto x = context.get_input(0); auto y = context.get_input(1); - auto adaptive_max_pool = context.mark_node(std::make_shared(x, y)); - auto return_indices = context.const_input(2); - OutputVector res{adaptive_max_pool->output(0)}; - if (return_indices) { - res.push_back(adaptive_max_pool->output(1)); - } - return res; + auto adaptive_max_pool = context.mark_node(std::make_shared(x, y, ov::element::i32)); + return {adaptive_max_pool->output(0), adaptive_max_pool->output(1)}; }; } // namespace op diff --git a/tests/layer_tests/pytorch_tests/test_adaptive_max_pool_2d.py b/tests/layer_tests/pytorch_tests/test_adaptive_max_pool_2d.py new file mode 100644 index 00000000000000..31030fb17a367e --- /dev/null +++ b/tests/layer_tests/pytorch_tests/test_adaptive_max_pool_2d.py @@ -0,0 +1,51 @@ +# Copyright (C) 2018-2022 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from pytorch_layer_test_class import PytorchLayerTest +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class TestAdaptiveMaxPool2D(PytorchLayerTest): + + def _prepare_input(self): + return (self.input_tensor,) + + def create_model(self, output_size=None, return_indices=False): + + class aten_adaptive_max_pool2d(torch.nn.Module): + + def __init__(self, output_size=None, return_indices=False) -> None: + super().__init__() + self.output_size = output_size + self.return_indices = return_indices + + def forward(self, input_tensor): + if self.return_indices: + output, indices = F.adaptive_max_pool2d(input_tensor, self.output_size, True) + return output + return F.adaptive_max_pool2d(input_tensor, self.output_size, False) + + ref_net = None + + return aten_adaptive_max_pool2d(output_size, return_indices), ref_net, "aten::adaptive_max_pool2d" + + @pytest.mark.parametrize('input_tensor', ([ + np.random.randn(1, 1, 4, 4).astype(np.float32), + np.random.randn(1, 3, 32, 32).astype(np.float32) + ])) + @pytest.mark.parametrize('output_size', ([ + [2, 2], + [4, 4], + ])) + @pytest.mark.parametrize('return_indices', ([ + False, + True, + ])) + @pytest.mark.nightly + def test_adaptive_max_pool2d(self, ie_device, precision, ir_version, input_tensor, output_size, return_indices): + self.input_tensor = input_tensor + self._test(*self.create_model(output_size, return_indices), ie_device, precision, ir_version)