Fix main element loading after create/update