Fixed the search_square function.