Fixed 'edit' functionality bug