diff --git a/src/frontends/pytorch/src/op/floor_divide.cpp b/src/frontends/pytorch/src/op/floor_divide.cpp index 77caa98e469..5731006dd77 100644 --- a/src/frontends/pytorch/src/op/floor_divide.cpp +++ b/src/frontends/pytorch/src/op/floor_divide.cpp @@ -18,6 +18,7 @@ OutputVector translate_floor_divide(NodeContext& context) { num_inputs_check(context, 2, 2); auto x = context.get_input(0); auto y = context.get_input(1); + align_eltwise_input_types(context, x, y, true); auto div = context.mark_node(std::make_shared(x, y, true)); return {context.mark_node(std::make_shared(div))}; }; diff --git a/src/frontends/pytorch/src/op/floordiv.cpp b/src/frontends/pytorch/src/op/floordiv.cpp index 48360a4a0ce..b85cacf3fc5 100644 --- a/src/frontends/pytorch/src/op/floordiv.cpp +++ b/src/frontends/pytorch/src/op/floordiv.cpp @@ -21,4 +21,4 @@ OutputVector translate_floordiv(NodeContext& context) { } // namespace op } // namespace pytorch } // namespace frontend -} // namespace ov \ No newline at end of file +} // namespace ov diff --git a/tests/layer_tests/pytorch_tests/test_floor_divide.py b/tests/layer_tests/pytorch_tests/test_floor_divide.py index 2d5592ee0e0..cd427acb3db 100644 --- a/tests/layer_tests/pytorch_tests/test_floor_divide.py +++ b/tests/layer_tests/pytorch_tests/test_floor_divide.py @@ -23,7 +23,20 @@ class TestFloorDivide(PytorchLayerTest): ref_net = None - # return aten_floor_divide(), ref_net, "aten::floor_divide" + return aten_floor_divide(), ref_net, "aten::floor_divide" + + def create_model_int(self): + import torch + + class aten_floor_divide(torch.nn.Module): + def __init__(self): + super(aten_floor_divide, self).__init__() + + def forward(self, input_tensor, other_tensor): + return torch.floor_divide(input_tensor.to(torch.int32), other_tensor.to(torch.int64)) + + ref_net = None + return aten_floor_divide(), ref_net, "aten::floor_divide" @pytest.mark.parametrize('input_tensor', ([ @@ -42,4 +55,23 @@ class TestFloorDivide(PytorchLayerTest): def test_floor_divide(self, input_tensor, other_tensor, ie_device, precision, ir_version): self.input_tensor = input_tensor self.other_tensor = other_tensor + self._test(*self.create_model(), ie_device, precision, ir_version, trace_model=True) + + @pytest.mark.parametrize('input_tensor', ([ + np.random.randint(low=0, high=10, size=5).astype(np.float32), + np.random.randint(low=1, high=10, size=(5, 5, 1)).astype(np.float32), + np.random.randint(low=1, high=10, size=(1, 1, 5, 5)).astype(np.float32), + ])) + @pytest.mark.parametrize('other_tensor', ([ + np.array([[2]]).astype(np.float32), + np.random.randint(low=1, high=10, size=5).astype(np.float32), + np.random.randint(low=1, high=10, size=(5, 1)).astype(np.float32), + np.random.randint(low=1, high=10, size=(1, 5)).astype(np.float32), + ])) + @pytest.mark.nightly + @pytest.mark.precommit + def test_floor_divide_int(self, input_tensor, other_tensor, ie_device, precision, ir_version): + self.input_tensor = input_tensor + self.other_tensor = other_tensor + self.create_model = self.create_model_int self._test(*self.create_model(), ie_device, precision, ir_version)