Проблема с вывод подменю в рубриках WP

Есть вот такой пример меню

С помощью данного кода я хочу вывести кроме основного меню еще и подменю в другом месте. Когда мы находимся на странице 1, 1.1, 1.3 или рубрике 1.2 должно отображаться подменю, тоже самое и с рубрикой 2.

Исходный код был взят из этого ответа. Я изменил немного структуру и подправил строку, что бы выводить подменю и на вложеных страницах/рубриках.

$current_item = ($item->menu_item_parent) ? $item->menu_item_parent : $item->ID;

Сейчас проблема в том, что данный код выводит подменю для страниц, но не работает с рубриками. В чем проблема и как ее исправить?

add_filter( 'wp_nav_menu_objects', 'wp_nav_menu_objects_filter', 10, 2 );
function wp_nav_menu_objects_filter( $sorted_menu_items, $args ) {
    if ( 'Main' !== $args->menu ) {
        return $sorted_menu_items;
    }
  $items = array();
  $current_item = null;

  foreach ( $sorted_menu_items  as $item) {
    if ( in_array( 'current_page_item', $item->classes, true ) ) {
      $current_item = ($item->menu_item_parent) ? $item->menu_item_parent : $item->ID;
        continue;
    }
  }
  foreach ( $sorted_menu_items  as $item) {
    if ( $current_item && intval( $item->menu_item_parent ) === intval( $current_item )) {
      $items[] = $item;
      continue;
    }
  }
    return $items;
}