Wordpress. Как заставить submenu быть внутри wrap родительского элемента при использовании walker?

Всем привет. Пытаюсь поменять структуру меню на сайте wordpress с помощью wp_nav_menu и walker. Структуру переделать получилось но, submenu вылазит за пределы родительского пункта меню. Подскажите пожалуйста как это можно исправить?
Вот сам вызоы wp_nav_menu:

$args = [
             'menu' => 'main_menu',
             'depth' => 2,
             'container' => 'nav',
             'container_class' => 'header_nav',
             'items_wrap' => '%3$s',
             'walker' => new Main_Menu_Walker()
];
wp_nav_menu($args);

А вот мой класс walker:

class Main_Menu_Walker extends Walker_Nav_Menu {
	/*
	 * Позволяет перезаписать <ul class="sub-menu">
	 */
	function start_lvl(&$output, $depth) {
		/*
		 * $depth – уровень вложенности, например 2,3 и т д
		 */
		$output .= '<div class="header_nav-link_wrap has-sublink">';
	}

	function end_lvl(&$output, $depth){
        $output .= "</div>";
    }


	/**
	 * @see Walker::start_el()
	 * @since 3.0.0
	 *
	 * @param string $output
	 * @param object $item Объект элемента меню, подробнее ниже.
	 * @param int $depth Уровень вложенности элемента меню.
	 * @param object $args Параметры функции wp_nav_menu
	 */
	function start_el(&$output, $item, $depth, $args) {
		global $wp_query;
		/*
		 * Генерируем строку с CSS-классами элемента меню
		 */
		if($depth == 0){
        	$class_names = 'header_nav-link_wrap';
    	}else{
    		$class_names = '';
    	}
        !empty ( $class_names ) and $class_names = ' class="'. esc_attr( $class_names ) . '"';
        if($depth == 0){
        	$output .= "<div $class_names>";
        }else{
        	$output .= "";
        }
        
        $attributes  = '';
        if($depth == 0){
        	$attributes .= 'class="header_nav-link"';
        }else{
        	$attributes .= 'class="header_nav-sublink"';
        }
        !empty( $item->url ) and $attributes .= ' href="'   . esc_attr( $item->url        ) .'"';
        $title = apply_filters( 'the_title', $item->title, $item->ID );
        if($depth == 0){
        	$item_output = $args->before
	        . "<a $attributes>"
	        . $args->link_before
	        . $title
	        . '</a></div>'
	        . $args->link_after
	        . $args->after;
        }else{
        	$item_output = $args->before
	        . "<a $attributes>"
	        . $args->link_before
	        . $title
	        . '</a>'
	        . $args->link_after
	        . $args->after;
        }
        
        $output .= apply_filters( 'walker_nav_menu_start_el', $item_output, $item, $depth, $args );
	}
}