Merge pull request #2 from sspintel/add-llama-ops

Add ops for LLAMA
cavusmustafa · Sep 14, 2023 · 430e4a0 · 430e4a0
2 parents f2eb270 + a7617bf
commit 430e4a0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 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
@@ -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
@@ -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},