From a7617bf474b5cca7b0cdc1004f7f81e282225b5e Mon Sep 17 00:00:00 2001
From: Suryaprakash Shanmugam <suryaprakash.shanmugam@intel.com>
Date: Mon, 11 Sep 2023 23:10:53 -0700
Subject: [PATCH] Add ops for LLAMA

---
 .../src/openvino/frontend/pytorch/torchdynamo/op_support.py   | 1 +
 src/frontends/pytorch/src/op_table.cpp                        | 4 +++-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/op_support.py b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/op_support.py
index 4cea76b9a98b50..b600e206c5207c 100644
--- a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/op_support.py
+++ b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/op_support.py
@@ -79,6 +79,7 @@ def __init__(self):
             "torch.ops.aten.logsumexp.default": None,
             #"torch.ops.aten.masked_fill.Scalar": None,
             "torch.ops.aten.masked_fill_.Scalar": None,
+            "torch.ops.aten.masked_fill.Scalar": None,
             "torch.ops.aten.masked_fill.Tensor": None,
             "torch.ops.aten.max_pool2d_with_indices.default": None,
             "torch.ops.aten.mean.dim": None,
diff --git a/src/frontends/pytorch/src/op_table.cpp b/src/frontends/pytorch/src/op_table.cpp
index cca45297eda451..7bc1c68a173f00 100644
--- a/src/frontends/pytorch/src/op_table.cpp
+++ b/src/frontends/pytorch/src/op_table.cpp
@@ -540,7 +540,8 @@ const std::map<std::string, CreatorFunction> get_supported_ops_fx() {
         {"aten.logical_not.default", op::translate_1to1_match_1_inputs<opset10::LogicalNot>},
         {"aten.logsumexp.default", op::translate_logsumexp},
         {"aten.lt.Tensor", op::translate_1to1_match_2_inputs_align_types<opset10::Less>},
-        {"aten.masked_fill.Scalar", op::translate_masked_fill},
+        {"aten.masked_fill_.Scalar", op::inplace_op<op::translate_masked_fill>},
+        {"aten.masked_fill.Scalar", op::inplace_op<op::translate_masked_fill>},
         {"aten.masked_fill.Tensor", op::translate_masked_fill},
         {"aten.masked_fill_.Scalar", op::inplace_op<op::translate_masked_fill>},
         {"aten.masked_fill_.Tensor", op::inplace_op<op::translate_masked_fill>},
@@ -573,6 +574,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_fx() {
         {"aten.split.Tensor", op::translate_chunk_fx},
         {"aten.split_with_sizes.default", op::translate_split_with_sizes_fx},
         {"aten.stack.default", op::translate_stack_fx},
+        {"aten.squeeze.dim", op::quantizable_op<op::translate_squeeze>},
         {"aten.sub.default", op::translate_sub},
         {"aten.sub.Tensor", op::translate_sub},
         {"aten.t.default", op::translate_t},