[PT FE] Fix issue with aten::index with bool mask on axis

src/frontends/pytorch/src/utils.cpp

@@ -797,8 +797,14 @@ bool index_tensor_on_list(ov::pass::NodeRegistry& rg,
         if (id_dtype == element::boolean || id_dtype == element::u8) {
             auto idx = rg.make<v0::Convert>(indices[i], element::u8);
             auto nonzero = rg.make<v3::NonZero>(idx);
-            auto input_order = rg.make<v0::Constant>(element::i32, Shape{2}, std::vector<int32_t>{1, 0});
-            auto masked_id = rg.make<v1::Transpose>(nonzero, input_order);
+            Output<Node> masked_id;
+            if (indices.size() == 1) {
+                auto input_order = rg.make<v0::Constant>(element::i32, Shape{2}, std::vector<int32_t>{1, 0});
+                masked_id = rg.make<v1::Transpose>(nonzero, input_order);
+            } else {
+                auto zero_const = rg.make<v0::Constant>(element::i32, Shape{1}, 0);
+                masked_id = rg.make<v0::Squeeze>(nonzero, zero_const);
+            }
             masked_indicies.push_back(masked_id);
             is_masked_bool.push_back(true);
         } else {
@@ -815,7 +821,7 @@ bool index_tensor_on_list(ov::pass::NodeRegistry& rg,
         return true;
     }
     // perform gather for single element case
-    if (advanced_ids.size() == 1) {
+    if (advanced_ids.size() == 1 && advanced_ids[0] == 0) {
         auto index = masked_indicies[advanced_ids[0]];
         if (is_masked_bool[advanced_ids[0]]) {
             auto gather = rg.make<v8::GatherND>(data, index);

tests/layer_tests/pytorch_tests/test_index.py

@@ -9,9 +9,11 @@
 
 
 class TestIndex(PytorchLayerTest):
-    def _prepare_input(self, input_shape, idx):
+    def _prepare_input(self, input_shape, idx=None):
         import numpy as np
-        return (np.random.randn(*input_shape).astype(np.float32), idx)
+        if idx is not None:
+            return (np.random.randn(*input_shape).astype(np.float32), idx)
+        return (np.random.randn(*input_shape).astype(np.float32),)
 
     def create_model(self, model="list"):
         import torch
@@ -35,11 +37,21 @@ class aten_index_getitem_bool(torch.nn.Module):
 
             def forward(self, x, idx):
                 return x.__getitem__(idx.to(torch.bool))
+
+        class aten_index_bool_with_axis(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.idx = torch.tensor([1, 0, 1, 0, 1], dtype=torch.bool)
+
+            def forward(self, x):
+                return x[:,:,self.idx]
+
         cases = {
             "list": aten_index_list,
             "getitem": aten_index_getitem,
             "list_with_bool": aten_index_list_bool,
-            "getitem_with_bool": aten_index_getitem_bool
+            "getitem_with_bool": aten_index_getitem_bool,
+            "bool_with_axis": aten_index_bool_with_axis
         }
 
         aten_index = cases[model]
@@ -74,6 +86,12 @@ def test_index_bool(self, input_shape, idx, case, ie_device, precision, ir_versi
         self._test(*self.create_model(case), ie_device, precision, ir_version,
                    kwargs_to_prepare_input={"input_shape": input_shape, "idx": idx})
 
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_index_bool_with_axis(self, ie_device, precision, ir_version):
+        self._test(*self.create_model("bool_with_axis"), ie_device, precision, ir_version,
+                   kwargs_to_prepare_input={"input_shape": (2, 2, 5)}, trace_model=True)
+
 
 class TestIndexRange(PytorchLayerTest):
     def _prepare_input(self, input_shape, idx):