Updated the update function