From 7b9db3d81b4bed45b3c2a9166428f849b44cff09 Mon Sep 17 00:00:00 2001
From: Maxim Vafin <maxim.vafin@intel.com>
Date: Mon, 30 Oct 2023 10:12:55 +0100
Subject: [PATCH] [PT FE] Add torch.int16 dtype support (#20735)

* Add torch.int16 dtype support

* Add test
---
 .../src/openvino/frontend/pytorch/utils.py    |  1 +
 .../py_frontend_tests/test_torch_decoder.py   | 22 +++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/src/bindings/python/src/openvino/frontend/pytorch/utils.py b/src/bindings/python/src/openvino/frontend/pytorch/utils.py
index a3ac46e7011..0901fea81c6 100644
--- a/src/bindings/python/src/openvino/frontend/pytorch/utils.py
+++ b/src/bindings/python/src/openvino/frontend/pytorch/utils.py
@@ -118,6 +118,7 @@ pt_to_ov_type_map = {
     "torch.float64": OVType.f64,
     "torch.uint8": OVType.u8,
     "torch.int8": OVType.i8,
+    "torch.int16": OVType.i16,
     "torch.int32": OVType.i32,
     "torch.int64": OVType.i64,
     "torch.bool": OVType.boolean,
diff --git a/tests/layer_tests/py_frontend_tests/test_torch_decoder.py b/tests/layer_tests/py_frontend_tests/test_torch_decoder.py
index 7e1758bd84d..dfacb8ecb85 100644
--- a/tests/layer_tests/py_frontend_tests/test_torch_decoder.py
+++ b/tests/layer_tests/py_frontend_tests/test_torch_decoder.py
@@ -249,6 +249,28 @@ def test_pytorch_decoder_can_convert_i8_tensor():
     assert ov_const[0].get_partial_shape() == PartialShape([2])
 
 
+@pytest.mark.precommit
+def test_pytorch_decoder_can_convert_i16_tensor():
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
+    from openvino.runtime import PartialShape, Type
+
+    class SomeTensor(torch.nn.Module):
+        def forward(self):
+            return torch.tensor([1, 2], dtype=torch.int16)
+
+    model = get_scripted_model(SomeTensor())
+    consts = [n for n in model.inlined_graph.nodes() if n.kind() ==
+              "prim::Constant"]
+    assert len(consts) > 0
+    some_const = consts[0]
+    nc_decoder = TorchScriptPythonDecoder(model, some_const)
+    ov_const = nc_decoder.as_constant()
+    assert ov_const is not None
+    assert len(ov_const) == 1
+    assert ov_const[0].get_element_type() == Type.i16
+    assert ov_const[0].get_partial_shape() == PartialShape([2])
+
+
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_i32_tensor():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder