add changes to fix remove button function